More type annotations

Author: maresb

barcode/__init__.py

@@ -5,8 +5,10 @@
 """
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 from typing import BinaryIO
+from typing import overload
 
 from barcode.codabar import CODABAR
 from barcode.codex import PZN
@@ -28,11 +30,10 @@
 from barcode.version import version  # noqa: F401
 
 if TYPE_CHECKING:
-    import os
-
+    from barcode.base import Barcode
     from barcode.writer import BaseWriter
 
-__BARCODE_MAP = {
+__BARCODE_MAP: dict[str, type[Barcode]] = {
     "codabar": CODABAR,
     "code128": Code128,
     "code39": Code39,
@@ -61,12 +62,29 @@
 PROVIDED_BARCODES.sort()
 
 
+@overload
+def get(
+    name: str, code: str, writer: BaseWriter | None = None, options: dict | None = None
+) -> Barcode:
+    ...
+
+
+@overload
+def get(
+    name: str,
+    code: None = None,
+    writer: BaseWriter | None = None,
+    options: dict | None = None,
+) -> type[Barcode]:
+    ...
+
+
 def get(
     name: str,
     code: str | None = None,
     writer: BaseWriter | None = None,
     options: dict | None = None,
-):
+) -> Barcode | type[Barcode]:
     """Helper method for getting a generator or even a generated code.
 
     :param name: The name of the type of barcode desired.
@@ -79,6 +97,7 @@ def get(
         generating.
     """
     options = options or {}
+    barcode: type[Barcode]
     try:
         barcode = __BARCODE_MAP[name.lower()]
     except KeyError as e:
@@ -89,15 +108,15 @@ def get(
     return barcode
 
 
-def get_class(name: str):
+def get_class(name: str) -> type[Barcode]:
     return get_barcode(name)
 
 
 def generate(
     name: str,
     code: str,
-    writer: BaseWriter | None = None,
-    output: str | (os.PathLike | (BinaryIO | None)) = None,
+    writer: BaseWriter | None,
+    output: str | os.PathLike | BinaryIO,
     writer_options: dict | None = None,
     text: str | None = None,
 ) -> str | None:
@@ -113,18 +132,22 @@ def generate(
     """
     from barcode.base import Barcode
 
+    if output is None:
+        raise TypeError("'output' cannot be None")
+
     writer = writer or Barcode.default_writer()
     writer.set_options(writer_options or {})
 
     barcode = get(name, code, writer)
 
     if isinstance(output, str):
         return barcode.save(output, writer_options, text)
-    if output:
-        barcode.write(output, writer_options, text)
+    if isinstance(output, os.PathLike):
+        with open(output, "wb") as fp:
+            barcode.write(fp, writer_options, text)
         return None
-
-    raise TypeError("'output' cannot be None")
+    barcode.write(output, writer_options, text)
+    return None
 
 
 get_barcode = get

barcode/base.py

@@ -34,6 +34,9 @@ class Barcode:
 
     writer: BaseWriter
 
+    def __init__(self, code: str, writer: BaseWriter | None = None, **options) -> None:
+        raise NotImplementedError
+
     def to_ascii(self) -> str:
         code_list = self.build()
         if not len(code_list) == 1:

barcode/writer.py

@@ -15,6 +15,9 @@
     from typing import Generator
     from typing import Literal
 
+    from PIL.Image import Image as T_Image
+    from PIL.ImageDraw import ImageDraw as T_ImageDraw
+
     class InternalText(TypedDict):
         start: list
         end: list
@@ -45,11 +48,11 @@ class Callbacks(TypedDict):
         Image = ImageDraw = ImageFont = None
 
 
-def mm2px(mm, dpi: int):
+def mm2px(mm: float, dpi: int) -> float:
     return (mm * dpi) / 25.4
 
 
-def pt2mm(pt):
+def pt2mm(pt: float) -> float:
     return pt * 0.352777778
 
 
@@ -58,8 +61,9 @@ def _set_attributes(element, **attributes):
         element.setAttribute(key, value)
 
 
-def create_svg_object(with_doctype=False):
+def create_svg_object(with_doctype=False) -> xml.dom.minidom.Document:
     imp = xml.dom.minidom.getDOMImplementation()
+    assert imp is not None
     doctype = imp.createDocumentType(
         "svg",
         "-//W3C//DTD SVG 1.1//EN",
@@ -317,13 +321,13 @@ def __init__(self) -> None:
             self._create_text,
             self._finish,
         )
-        self.compress = False
-        self.with_doctype = True
-        self._document = None
-        self._root = None
-        self._group = None
+        self.compress: bool = False
+        self.with_doctype: bool = True
+        self._document: xml.dom.minidom.Document
+        self._root: xml.dom.minidom.Element
+        self._group: xml.dom.minidom.Element
 
-    def _init(self, code):
+    def _init(self, code: list[str]):
         width, height = self.calculate_size(len(code[0]), len(code))
         self._document = create_svg_object(self.with_doctype)
         self._root = self._document.documentElement
@@ -438,16 +442,18 @@ def __init__(self, format="PNG", mode="RGB", dpi=300) -> None:
             self.format = format
             self.mode = mode
             self.dpi = dpi
-            self._image = None
-            self._draw = None
+            self._image: T_Image
+            self._draw: T_ImageDraw
 
-        def _init(self, code):
+        def _init(self, code: list[str]) -> None:
+            if ImageDraw is None:
+                raise RuntimeError("Pillow not found. Cannot create image.")
             width, height = self.calculate_size(len(code[0]), len(code))
             size = (int(mm2px(width, self.dpi)), int(mm2px(height, self.dpi)))
             self._image = Image.new(self.mode, size, self.background)
             self._draw = ImageDraw.Draw(self._image)
 
-        def _paint_module(self, xpos, ypos, width, color):
+        def _paint_module(self, xpos: float, ypos: float, width: float, color):
             size = [
                 (mm2px(xpos, self.dpi), mm2px(ypos, self.dpi)),
                 (
@@ -458,6 +464,7 @@ def _paint_module(self, xpos, ypos, width, color):
             self._draw.rectangle(size, outline=color, fill=color)
 
         def _paint_text(self, xpos, ypos):
+            assert ImageFont is not None
             font_size = int(mm2px(pt2mm(self.font_size), self.dpi))
             if font_size <= 0:
                 return
@@ -472,7 +479,7 @@ def _paint_text(self, xpos, ypos):
                 )
                 ypos += pt2mm(self.font_size) / 2 + self.text_line_distance
 
-        def _finish(self) -> Image:
+        def _finish(self) -> T_Image:
             return self._image
 
         def save(self, filename: str, output) -> str: