Add type hints to all code

src/rapids_pre_commit_hooks/alpha_spec.py

@@ -16,34 +16,35 @@
 import os
 import re
 from functools import cache, total_ordering
+from typing import Optional
 
 import yaml
 from packaging.requirements import InvalidRequirement, Requirement
-from rapids_metadata.metadata import RAPIDSVersion
+from rapids_metadata.metadata import RAPIDSMetadata, RAPIDSVersion
 from rapids_metadata.remote import fetch_latest
 
-from .lint import LintMain
+from .lint import Linter, LintMain
 
-ALPHA_SPECIFIER = ">=0.0.0a0"
+ALPHA_SPECIFIER: str = ">=0.0.0a0"
 
-ALPHA_SPEC_OUTPUT_TYPES = {
+ALPHA_SPEC_OUTPUT_TYPES: set[str] = {
     "pyproject",
     "requirements",
 }
 
-CUDA_SUFFIX_REGEX = re.compile(r"^(?P<package>.*)-cu[0-9]{2}$")
+CUDA_SUFFIX_REGEX: re.Pattern = re.compile(r"^(?P<package>.*)-cu[0-9]{2}$")
 
 
 @cache
-def all_metadata():
+def all_metadata() -> "RAPIDSMetadata":
     return fetch_latest()
 
 
-def node_has_type(node, tag_type):
+def node_has_type(node: "yaml.Node", tag_type: str) -> bool:
     return node.tag == f"tag:yaml.org,2002:{tag_type}"
 
 
-def get_rapids_version(args: argparse.Namespace) -> RAPIDSVersion:
+def get_rapids_version(args: argparse.Namespace) -> "RAPIDSVersion":
     md = all_metadata()
     return (
         md.versions[args.rapids_version]
@@ -60,7 +61,9 @@ def strip_cuda_suffix(args: argparse.Namespace, name: str) -> str:
     return name
 
 
-def check_and_mark_anchor(anchors, used_anchors, node):
+def check_and_mark_anchor(
+    anchors: dict[str, "yaml.Node"], used_anchors: set[str], node: "yaml.Node"
+) -> tuple[bool, Optional[str]]:
     for key, value in anchors.items():
         if value == node:
             anchor = key
@@ -74,16 +77,24 @@ def check_and_mark_anchor(anchors, used_anchors, node):
     return True, anchor
 
 
-def check_package_spec(linter, args, anchors, used_anchors, node):
+def check_package_spec(
+    linter: Linter,
+    args: argparse.Namespace,
+    anchors: dict[str, "yaml.Node"],
+    used_anchors: set[str],
+    node: "yaml.Node",
+) -> None:
     @total_ordering
     class SpecPriority:
-        def __init__(self, spec):
-            self.spec = spec
+        def __init__(self, spec: str):
+            self.spec: str = spec
 
-        def __eq__(self, other):
+        def __eq__(self, other: object) -> bool:
+            assert isinstance(other, SpecPriority)
             return self.spec == other.spec
 
-        def __lt__(self, other):
+        def __lt__(self, other: object) -> bool:
+            assert isinstance(other, SpecPriority)
             if self.spec == other.spec:
                 return False
             if self.spec == ALPHA_SPECIFIER:
@@ -92,10 +103,10 @@ def __lt__(self, other):
                 return True
             return self.sort_str() < other.sort_str()
 
-        def sort_str(self):
+        def sort_str(self) -> str:
             return "".join(c for c in self.spec if c not in "<>=")
 
-    def create_specifier_string(specifiers):
+    def create_specifier_string(specifiers: set[str]) -> str:
         return ",".join(sorted(specifiers, key=SpecPriority))
 
     if node_has_type(node, "str"):
@@ -140,15 +151,27 @@ def create_specifier_string(specifiers):
                     )
 
 
-def check_packages(linter, args, anchors, used_anchors, node):
+def check_packages(
+    linter: Linter,
+    args: argparse.Namespace,
+    anchors: dict[str, "yaml.Node"],
+    used_anchors: set[str],
+    node: "yaml.Node",
+) -> None:
     if node_has_type(node, "seq"):
         descend, _ = check_and_mark_anchor(anchors, used_anchors, node)
         if descend:
             for package_spec in node.value:
                 check_package_spec(linter, args, anchors, used_anchors, package_spec)
 
 
-def check_common(linter, args, anchors, used_anchors, node):
+def check_common(
+    linter: Linter,
+    args: argparse.Namespace,
+    anchors: dict[str, "yaml.Node"],
+    used_anchors: set[str],
+    node: "yaml.Node",
+) -> None:
     if node_has_type(node, "seq"):
         for dependency_set in node.value:
             if node_has_type(dependency_set, "map"):
@@ -162,7 +185,13 @@ def check_common(linter, args, anchors, used_anchors, node):
                         )
 
 
-def check_matrices(linter, args, anchors, used_anchors, node):
+def check_matrices(
+    linter: Linter,
+    args: argparse.Namespace,
+    anchors: dict[str, "yaml.Node"],
+    used_anchors: set[str],
+    node: "yaml.Node",
+) -> None:
     if node_has_type(node, "seq"):
         for item in node.value:
             if node_has_type(item, "map"):
@@ -176,7 +205,13 @@ def check_matrices(linter, args, anchors, used_anchors, node):
                         )
 
 
-def check_specific(linter, args, anchors, used_anchors, node):
+def check_specific(
+    linter: Linter,
+    args: argparse.Namespace,
+    anchors: dict[str, "yaml.Node"],
+    used_anchors: set[str],
+    node: "yaml.Node",
+) -> None:
     if node_has_type(node, "seq"):
         for matrix_matcher in node.value:
             if node_has_type(matrix_matcher, "map"):
@@ -190,7 +225,13 @@ def check_specific(linter, args, anchors, used_anchors, node):
                         )
 
 
-def check_dependencies(linter, args, anchors, used_anchors, node):
+def check_dependencies(
+    linter: Linter,
+    args: argparse.Namespace,
+    anchors: dict[str, "yaml.Node"],
+    used_anchors: set[str],
+    node: "yaml.Node",
+) -> None:
     if node_has_type(node, "map"):
         for _, dependencies_value in node.value:
             if node_has_type(dependencies_value, "map"):
@@ -206,7 +247,13 @@ def check_dependencies(linter, args, anchors, used_anchors, node):
                             )
 
 
-def check_root(linter, args, anchors, used_anchors, node):
+def check_root(
+    linter: Linter,
+    args: argparse.Namespace,
+    anchors: dict[str, "yaml.Node"],
+    used_anchors: set[str],
+    node: "yaml.Node",
+) -> None:
     if node_has_type(node, "map"):
         for root_key, root_value in node.value:
             if node_has_type(root_key, "str") and root_key.value == "dependencies":
@@ -219,35 +266,37 @@ class AnchorPreservingLoader(yaml.SafeLoader):
     dictionary for each parsed document.
     """
 
-    def __init__(self, stream):
+    def __init__(self, stream) -> None:
         super().__init__(stream)
-        self.document_anchors = []
+        self.document_anchors: list[dict[str, yaml.Node]] = []
 
-    def compose_document(self):
+    def compose_document(self) -> "yaml.Node":
         # Drop the DOCUMENT-START event.
         self.get_event()
 
         # Compose the root node.
-        node = self.compose_node(None, None)
+        node = self.compose_node(None, None)  # type: ignore[arg-type]
 
         # Drop the DOCUMENT-END event.
         self.get_event()
 
         self.document_anchors.append(self.anchors)
         self.anchors = {}
+        assert node is not None
         return node
 
 
-def check_alpha_spec(linter, args):
+def check_alpha_spec(linter: Linter, args: argparse.Namespace) -> None:
     loader = AnchorPreservingLoader(linter.content)
     try:
         root = loader.get_single_node()
+        assert root is not None
     finally:
         loader.dispose()
     check_root(linter, args, loader.document_anchors[0], set(), root)
 
 
-def main():
+def main() -> None:
     m = LintMain()
     m.argparser.description = (
         "Verify that RAPIDS packages in dependencies.yaml do (or do not) have "