Attempt to type-check writers

Author: WhyNotHugo

barcode/base.py

@@ -4,15 +4,20 @@
 
 from typing import TYPE_CHECKING
 from typing import ClassVar
+from typing import Generic
+from typing import TypeVar
 
 from barcode.writer import BaseWriter
 from barcode.writer import SVGWriter
+from barcode.writer import T_Output
 
 if TYPE_CHECKING:
     from typing import BinaryIO
 
+W = TypeVar("W", bound=BaseWriter[object])
 
-class Barcode:
+
+class Barcode(Generic[W, T_Output]):
     name = ""
 
     digits = 0
@@ -31,9 +36,9 @@ class Barcode:
         "text": "",
     }
 
-    writer: BaseWriter
+    writer: W
 
-    def __init__(self, code: str, writer: BaseWriter | None = None, **options) -> None:
+    def __init__(self, code: str, writer: W | None = None, **options) -> None:
         raise NotImplementedError
 
     def to_ascii(self) -> str:
@@ -62,7 +67,10 @@ def get_fullcode(self):
         raise NotImplementedError
 
     def save(
-        self, filename: str, options: dict | None = None, text: str | None = None
+        self,
+        filename: str,
+        options: dict | None = None,
+        text: str | None = None,
     ) -> str:
         """Renders the barcode and saves it in `filename`.
 
@@ -72,7 +80,7 @@ def save(
 
         :returns: The full filename with extension.
         """
-        output = self.render(options, text) if text else self.render(options)
+        output: T_Output = self.render(options, text) if text else self.render(options)
 
         return self.writer.save(filename, output)
 
@@ -92,7 +100,11 @@ def write(
         output = self.render(options, text)
         self.writer.write(output, fp)
 
-    def render(self, writer_options: dict | None = None, text: str | None = None):
+    def render(
+        self,
+        writer_options: dict | None = None,
+        text: str | None = None,
+    ) -> T_Output:
         """Renders the barcode using `self.writer`.
 
         :param writer_options: Options for `self.writer`, see writer docs for details.

barcode/writer.py

@@ -3,10 +3,14 @@
 import gzip
 import os
 import xml.dom.minidom
+from abc import ABC
+from abc import abstractmethod
 from typing import TYPE_CHECKING
 from typing import BinaryIO
 from typing import Callable
+from typing import Generic
 from typing import TypedDict
+from typing import TypeVar
 
 from barcode.version import version
 
@@ -78,8 +82,10 @@ def create_svg_object(with_doctype: bool = False) -> xml.dom.minidom.Document:
 COMMENT = f"Autogenerated with python-barcode {version}"
 PATH = os.path.dirname(os.path.abspath(__file__))
 
+T_Output = TypeVar("T_Output")
 
-class BaseWriter:
+
+class BaseWriter(ABC, Generic[T_Output]):
     """Baseclass for all writers.
 
     Initializes the basic writer options. Child classes can add more attributes and can
@@ -164,7 +170,8 @@ def calculate_size(self, modules_per_line: int, number_of_lines: int) -> tuple:
             height += self.text_line_distance * (number_of_text_lines - 1)
         return width, height
 
-    def save(self, filename: str, output) -> str:
+    @abstractmethod
+    def save(self, filename: str, output: T_Output) -> str:
         """Saves the rendered output to `filename`.
 
         :param filename: Filename without extension.
@@ -307,11 +314,12 @@ def render(self, code: list[str]):
 
         return self._callbacks["finish"]()
 
+    @abstractmethod
     def write(self, content, fp: BinaryIO) -> None:
         raise NotImplementedError
 
 
-class SVGWriter(BaseWriter):
+class SVGWriter(BaseWriter[bytes]):
     def __init__(self) -> None:
         super().__init__(
             self._init,
@@ -398,7 +406,7 @@ def _finish(self) -> bytes:
             indent=4 * " ", newl=os.linesep, encoding="UTF-8"
         )
 
-    def save(self, filename: str, output) -> str:
+    def save(self, filename: str, output: bytes) -> str:
         if self.compress:
             _filename = f"{filename}.svgz"
             with gzip.open(_filename, "wb") as f:
@@ -419,7 +427,21 @@ def write(self, content, fp: BinaryIO) -> None:
 
 
 if Image is None:
-    ImageWriter: type | None = None
+    if TYPE_CHECKING:
+
+        class ImageWriter(BaseWriter):
+            def __init__(
+                self,
+                format: str = "PNG",
+                mode: str = "RGB",
+                dpi: int = 300,
+            ) -> None: ...
+
+            def save(self, filename: str, output: T_Image) -> str: ...
+
+            def write(self, content, fp: BinaryIO) -> None: ...
+    else:
+        ImageWriter = None
 else:
 
     class ImageWriter(BaseWriter):  # type: ignore[no-redef]
@@ -497,7 +519,7 @@ def _paint_text(self, xpos, ypos):
         def _finish(self) -> T_Image:
             return self._image
 
-        def save(self, filename: str, output) -> str:
+        def save(self, filename: str, output: T_Image) -> str:
             filename = f"{filename}.{self.format.lower()}"
             output.save(filename, self.format.upper())
             return filename

tests/test_writers.py

@@ -5,12 +5,13 @@
 
 from barcode import EAN13
 from barcode.writer import ImageWriter
+from barcode.writer import Image
 from barcode.writer import SVGWriter
 
 PATH = os.path.dirname(os.path.abspath(__file__))
 TESTPATH = os.path.join(PATH, "test_outputs")
 
-if ImageWriter is not None:
+if Image is not None:
 
     def test_saving_image_to_byteio() -> None:
         assert ImageWriter is not None  # workaround for mypy