Add missing annotations

Author: insolor

dawg_python/dawgs.py

@@ -26,7 +26,7 @@ class DAWG:
     def __init__(self) -> None:
         self.dct = None
 
-    def __contains__(self, key) -> bool:
+    def __contains__(self, key: str | bytes) -> bool:
         if not isinstance(key, bytes):
             key = key.encode("utf8")
         return self.dct.contains(key)
@@ -73,7 +73,7 @@ def _similar_keys(self, current_prefix: str, key: str, index: int, replace_chars
 
         return res
 
-    def similar_keys(self, key: str, replaces: CompiledReplaces):
+    def similar_keys(self, key: str, replaces: CompiledReplaces) -> list[str]:
         """
         Returns all variants of ``key`` in this DAWG according to
         ``replaces``.
@@ -202,7 +202,7 @@ def __contains__(self, key: str | bytes) -> bool:
             key = key.encode("utf8")
         return bool(self._follow_key(key))
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str | bytes) -> list[bytes]:
         res = self.get(key)
         if res is None:
             raise KeyError(key)
@@ -241,7 +241,7 @@ def _value_for_index(self, index: int) -> list[bytes]:
 
         return res
 
-    def b_get_value(self, b_key) -> list[bytes]:
+    def b_get_value(self, b_key: bytes) -> list[bytes]:
         index = self._follow_key(b_key)
         if not index:
             return []

dawg_python/units.py

@@ -10,17 +10,17 @@
 EXTENSION_BIT = 1 << 9
 
 
-def has_leaf(base: int, _mask=HAS_LEAF_BIT) -> bool:
+def has_leaf(base: int, _mask: int = HAS_LEAF_BIT) -> bool:
     """Check if a unit has a leaf as a child or not."""
     return bool(base & _mask)
 
 
-def value(base: int, _mask=~IS_LEAF_BIT & PRECISION_MASK) -> int:
+def value(base: int, _mask: int = ~IS_LEAF_BIT & PRECISION_MASK) -> int:
     """Check if a unit corresponds to a leaf or not."""
     return base & _mask
 
 
-def label(base: int, _mask=IS_LEAF_BIT | 0xFF) -> int:
+def label(base: int, _mask: int = IS_LEAF_BIT | 0xFF) -> int:
     """Read a label with a leaf flag from a non-leaf unit."""
     return base & _mask
 

dawg_python/wrapper.py

@@ -8,6 +8,7 @@
 
 if TYPE_CHECKING:
     from io import BytesIO
+    from pathlib import Path
 
     from typing_extensions import Self
 
@@ -74,7 +75,7 @@ def follow_bytes(self, s: bytes, index: int) -> int | None:
         return index
 
     @classmethod
-    def load(cls, path) -> Self:
+    def load(cls, path: str | Path) -> Self:
         dawg = cls()
         with open(path, "rb") as f:
             dawg.read(f)

pyproject.toml

@@ -42,21 +42,17 @@ target-version = "py38"
 [tool.ruff.lint]
 select = ["ALL"]
 ignore = [
-    "ANN001",
-    "ANN002",
-    "ANN201",
-    "ANN202",
-    "ANN204",
-    "ANN206",
-    # ^^^ Silence warnings about the type annotations for now
-
     "D",
     "PTH123",
 ]
 fixable = ["ALL"]
 
 [tool.ruff.lint.per-file-ignores]
 "_prepare_dev_data.py" = [
+    "ANN001",
+    "ANN002",
+    "ANN201",
+    "ANN202",
     "ERA001",
 ]
 
@@ -65,6 +61,10 @@ fixable = ["ALL"]
 ]
 
 "bench/*" = [
+    "ANN001",
+    "ANN002",
+    "ANN201",
+    "ANN202",
     "ERA001",
     "PERF401",
     "PLR0913",
@@ -83,4 +83,4 @@ fixable = ["ALL"]
     "S101",
     "S301",
     "S324",
-]
+]