Add type hints for the resolver module

.pre-commit-config.yaml

@@ -20,7 +20,7 @@ repos:
       additional_dependencies:
       - flake8-pytest-style
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.790
+    rev: v0.800
     hooks:
       - id: mypy
         # Avoid error: Duplicate module named 'setup'

piptools/cache.py

@@ -2,7 +2,9 @@
 import os
 import platform
 import sys
+from typing import Dict, Sequence, Set, Tuple
 
+from pip._internal.req.req_install import InstallRequirement
 from pip._vendor.packaging.requirements import Requirement
 
 from .exceptions import PipToolsError
@@ -76,7 +78,7 @@ def cache(self):
             self.read_cache()
         return self._cache
 
-    def as_cache_key(self, ireq):
+    def as_cache_key(self, ireq: InstallRequirement) -> Tuple[str, str]:
         """
         Given a requirement, return its cache key. This behavior is a little weird
         in order to allow backwards compatibility with cache files. For a requirement
@@ -103,13 +105,13 @@ def read_cache(self):
         except FileNotFoundError:
             self._cache = {}
 
-    def write_cache(self):
+    def write_cache(self) -> None:
         """Writes the cache to disk as JSON."""
         doc = {"__format__": 1, "dependencies": self._cache}
         with open(self._cache_file, "w") as f:
             json.dump(doc, f, sort_keys=True)
 
-    def clear(self):
+    def clear(self) -> None:
         self._cache = {}
         self.write_cache()
 
@@ -127,7 +129,9 @@ def __setitem__(self, ireq, values):
         self.cache[pkgname][pkgversion_and_extras] = values
         self.write_cache()
 
-    def reverse_dependencies(self, ireqs):
+    def reverse_dependencies(
+        self, ireqs: Sequence[InstallRequirement]
+    ) -> Dict[str, Set[str]]:
         """
         Returns a lookup table of reverse dependencies for all the given ireqs.
 
@@ -139,7 +143,9 @@ def reverse_dependencies(self, ireqs):
         ireqs_as_cache_values = [self.as_cache_key(ireq) for ireq in ireqs]
         return self._reverse_dependencies(ireqs_as_cache_values)
 
-    def _reverse_dependencies(self, cache_keys):
+    def _reverse_dependencies(
+        self, cache_keys: Sequence[Tuple[str, str]]
+    ) -> Dict[str, Set[str]]:
         """
         Returns a lookup table of reverse dependencies for all the given cache keys.
 

piptools/logging.py

@@ -23,7 +23,7 @@ def log(self, message: str, *args: Any, **kwargs: Any) -> None:
         prefix = " " * self.current_indent
         click.secho(prefix + message, *args, **kwargs)
 
-    def debug(self, *args, **kwargs):
+    def debug(self, *args: Any, **kwargs: Any) -> None:
         if self.verbosity >= 1:
             self.log(*args, **kwargs)
 
@@ -35,14 +35,14 @@ def warning(self, *args: Any, **kwargs: Any) -> None:
         kwargs.setdefault("fg", "yellow")
         self.log(*args, **kwargs)
 
-    def error(self, *args, **kwargs):
+    def error(self, *args: Any, **kwargs: Any) -> None:
         kwargs.setdefault("fg", "red")
         self.log(*args, **kwargs)
 
-    def _indent(self):
+    def _indent(self) -> None:
         self.current_indent += self._indent_width
 
-    def _dedent(self):
+    def _dedent(self) -> None:
         self.current_indent -= self._indent_width
 
     @contextlib.contextmanager

piptools/repositories/base.py

@@ -1,9 +1,12 @@
 from abc import ABCMeta, abstractmethod
 from contextlib import contextmanager
+from typing import Optional, Set
+
+from pip._internal.req.req_install import InstallRequirement
 
 
 class BaseRepository(metaclass=ABCMeta):
-    def clear_caches(self):
+    def clear_caches(self) -> None:
         """Should clear any caches used by the implementation."""
 
     @abstractmethod
@@ -12,24 +15,26 @@ def freshen_build_caches(self):
         """Should start with fresh build/source caches."""
 
     @abstractmethod
-    def find_best_match(self, ireq):
+    def find_best_match(
+        self, ireq: InstallRequirement, prereleases: Optional[bool]
+    ) -> InstallRequirement:
         """
         Return a Version object that indicates the best match for the given
         InstallRequirement according to the repository.
         """
 
     @abstractmethod
-    def get_dependencies(self, ireq):
+    def get_dependencies(self, ireq: InstallRequirement) -> Set[InstallRequirement]:
         """
         Given a pinned, URL, or editable InstallRequirement, returns a set of
         dependencies (also InstallRequirements, but not necessarily pinned).
         They indicate the secondary dependencies for the given requirement.
         """
 
     @abstractmethod
-    def get_hashes(self, ireq):
+    def get_hashes(self, ireq: InstallRequirement) -> Set[str]:
         """
-        Given a pinned InstallRequire, returns a set of hashes that represent
+        Given a pinned InstallRequirement, returns a set of hashes that represent
         all of the files for a given requirement. It is not acceptable for an
         editable or unpinned requirement to be passed to this function.
         """
@@ -42,7 +47,9 @@ def allow_all_wheels(self):
         """
 
     @abstractmethod
-    def copy_ireq_dependencies(self, source, dest):
+    def copy_ireq_dependencies(
+        self, source: InstallRequirement, dest: InstallRequirement
+    ) -> None:
         """
         Notifies the repository that `dest` is a copy of `source`, and so it
         has the same dependencies. Otherwise, once we prepare an ireq to assign

piptools/resolver.py

@@ -1,11 +1,16 @@
 import copy
 from functools import partial
 from itertools import chain, count, groupby
+from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple
 
 import click
 from pip._internal.req.constructors import install_req_from_line
+from pip._internal.req.req_install import InstallRequirement
 from pip._internal.req.req_tracker import update_env_context_manager
 
+from piptools.cache import DependencyCache
+from piptools.repositories.base import BaseRepository
+
 from .logging import log
 from .utils import (
     UNSAFE_PACKAGES,
@@ -25,27 +30,31 @@ class RequirementSummary:
     Summary of a requirement's properties for comparison purposes.
     """
 
-    def __init__(self, ireq):
+    def __init__(self, ireq: InstallRequirement) -> None:
         self.req = ireq.req
         self.key = key_from_ireq(ireq)
         self.extras = frozenset(ireq.extras)
         self.specifier = ireq.specifier
 
-    def __eq__(self, other):
-        return (
-            self.key == other.key
-            and self.specifier == other.specifier
-            and self.extras == other.extras
-        )
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, self.__class__):
+            return (
+                self.key == other.key
+                and self.specifier == other.specifier
+                and self.extras == other.extras
+            )
+        return NotImplemented
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash((self.key, self.specifier, self.extras))
 
-    def __str__(self):
+    def __str__(self) -> str:
         return repr((self.key, str(self.specifier), sorted(self.extras)))
 
 
-def combine_install_requirements(repository, ireqs):
+def combine_install_requirements(
+    repository: BaseRepository, ireqs: Iterable[InstallRequirement]
+) -> InstallRequirement:
     """
     Return a single install requirement that reflects a combination of
     all the inputs.
@@ -100,34 +109,36 @@ def combine_install_requirements(repository, ireqs):
 class Resolver:
     def __init__(
         self,
-        constraints,
-        repository,
-        cache,
-        prereleases=False,
-        clear_caches=False,
-        allow_unsafe=False,
-    ):
+        constraints: List[InstallRequirement],
+        repository: BaseRepository,
+        cache: DependencyCache,
+        prereleases: Optional[bool] = False,
+        clear_caches: bool = False,
+        allow_unsafe: bool = False,
+    ) -> None:
         """
         This class resolves a given set of constraints (a collection of
         InstallRequirement objects) by consulting the given Repository and the
         DependencyCache.
         """
         self.our_constraints = set(constraints)
-        self.their_constraints = set()
+        self.their_constraints: Set[InstallRequirement] = set()
         self.repository = repository
         self.dependency_cache = cache
         self.prereleases = prereleases
         self.clear_caches = clear_caches
         self.allow_unsafe = allow_unsafe
-        self.unsafe_constraints = set()
+        self.unsafe_constraints: Set[InstallRequirement] = set()
 
     @property
-    def constraints(self):
+    def constraints(self) -> Set[InstallRequirement]:
         return set(
             self._group_constraints(chain(self.our_constraints, self.their_constraints))
         )
 
-    def resolve_hashes(self, ireqs):
+    def resolve_hashes(
+        self, ireqs: Set[InstallRequirement]
+    ) -> Dict[InstallRequirement, Set[str]]:
         """
         Finds acceptable hashes for all of the given InstallRequirements.
         """
@@ -136,7 +147,7 @@ def resolve_hashes(self, ireqs):
         with self.repository.allow_all_wheels(), log.indentation():
             return {ireq: self.repository.get_hashes(ireq) for ireq in ireqs}
 
-    def resolve(self, max_rounds=10):
+    def resolve(self, max_rounds: int = 10) -> Set[InstallRequirement]:
         """
         Finds concrete package versions for all the given InstallRequirements
         and their recursive dependencies.  The end result is a flat list of
@@ -193,7 +204,7 @@ def resolve(self, max_rounds=10):
             # sense for installation tools) so this seems sufficient.
             reverse_dependencies = self.reverse_dependencies(results)
             for req in results.copy():
-                required_by = reverse_dependencies.get(req.name.lower(), [])
+                required_by = reverse_dependencies.get(req.name.lower(), set())
                 if req.name in UNSAFE_PACKAGES or (
                     required_by and all(name in UNSAFE_PACKAGES for name in required_by)
                 ):
@@ -202,7 +213,9 @@ def resolve(self, max_rounds=10):
 
         return results
 
-    def _group_constraints(self, constraints):
+    def _group_constraints(
+        self, constraints: Iterable[InstallRequirement]
+    ) -> Iterator[InstallRequirement]:
         """
         Groups constraints (remember, InstallRequirements!) by their key name,
         and combining their SpecifierSets into a single InstallRequirement per
@@ -235,7 +248,7 @@ def _group_constraints(self, constraints):
         ):
             yield combine_install_requirements(self.repository, ireqs)
 
-    def _resolve_one_round(self):
+    def _resolve_one_round(self) -> Tuple[bool, Set[InstallRequirement]]:
         """
         Resolves one level of the current constraints, by finding the best
         match for each package in the repository and adding all requirements
@@ -263,7 +276,7 @@ def _resolve_one_round(self):
         log.debug("")
         log.debug("Finding secondary dependencies:")
 
-        their_constraints = []
+        their_constraints: List[InstallRequirement] = []
         with log.indentation():
             for best_match in best_matches:
                 their_constraints.extend(self._iter_dependencies(best_match))
@@ -295,7 +308,7 @@ def _resolve_one_round(self):
         self.their_constraints = theirs
         return has_changed, best_matches
 
-    def get_best_match(self, ireq):
+    def get_best_match(self, ireq: InstallRequirement) -> InstallRequirement:
         """
         Returns a (pinned or editable) InstallRequirement, indicating the best
         match to use for the given InstallRequirement (in the form of an
@@ -338,7 +351,9 @@ def get_best_match(self, ireq):
             best_match._source_ireqs = ireq._source_ireqs
         return best_match
 
-    def _iter_dependencies(self, ireq):
+    def _iter_dependencies(
+        self, ireq: InstallRequirement
+    ) -> Iterator[InstallRequirement]:
         """
         Given a pinned, url, or editable InstallRequirement, collects all the
         secondary dependencies for them, either by looking them up in a local
@@ -389,7 +404,9 @@ def _iter_dependencies(self, ireq):
                 dependency_string, constraint=ireq.constraint, comes_from=ireq
             )
 
-    def reverse_dependencies(self, ireqs):
+    def reverse_dependencies(
+        self, ireqs: Set[InstallRequirement]
+    ) -> Dict[str, Set[str]]:
         non_editable = [
             ireq for ireq in ireqs if not (ireq.editable or is_url_requirement(ireq))
         ]