perf: cache comparison key for choosing the next package to resolve

Author: radoering

src/poetry/mixology/version_solver.py

@@ -4,6 +4,7 @@
 import functools
 import time
 
+from enum import IntEnum
 from typing import TYPE_CHECKING
 from typing import Optional
 
@@ -31,6 +32,22 @@
 _conflict = object()
 
 
+class Preference(IntEnum):
+    """
+    Preference is one of the criteria for choosing which dependency to solve
+    first. A higher value means that there are "more options" to satisfy
+    a dependency. A lower value takes precedence.
+    """
+
+    DIRECT_ORIGIN = 0
+    NO_CHOICE = 1
+    USE_LATEST = 2
+    LOCKED = 3
+    DEFAULT = 4
+
+
+CompKey = tuple[Preference, int, bool, int]
+
 DependencyCacheKey = tuple[
     str, Optional[str], Optional[str], Optional[str], Optional[str]
 ]
@@ -149,6 +166,7 @@ def __init__(self, root: ProjectPackage, provider: Provider) -> None:
             int, set[Incompatibility]
         ] = collections.defaultdict(set)
         self._solution = PartialSolution()
+        self._get_comp_key_cached = functools.cache(self._get_comp_key)
 
     @property
     def solution(self) -> PartialSolution:
@@ -432,107 +450,111 @@ def _resolve_conflict(self, incompatibility: Incompatibility) -> Incompatibility
 
         raise SolveFailureError(incompatibility)
 
-    def _choose_next(self, unsatisfied: list[Dependency]) -> Dependency:
+    def _get_comp_key(self, dependency: Dependency) -> CompKey:
         """
+        Returns a tuple of
+        - preference
+        - num_deps_upper_bound
+        - has_deps
+        - num_packages
+        that serves as priority for choosing the next package to resolve.
+        (A lower value takes precedence.)
+
+        In order to provide results that are as deterministic as possible
+        and consistent between `poetry lock` and `poetry update`, the return value
+        of two different dependencies should not be equal if possible.
+
+        ## preference
+
+        See Preference class.
+
+        ## num_deps_upper_bound
+
+        A dependency with an upper bound is more likely to cause conflicts. Therefore,
+        a package with more dependencies with upper bounds should be chosen first.
+
+        ## has_deps
+
+        A package with dependencies should be chosen first
+        because a package without dependencies is less likely to cause conflicts.
+
+        ## num_packages
+
         The original algorithm proposes to prefer packages with as few remaining
         versions as possible, so that if a conflict is necessary it's forced quickly.
         https://github.com/dart-lang/pub/blob/master/doc/solver.md#decision-making
         However, this leads to the famous boto3 vs. urllib3 issue, so we prefer
         packages with more remaining versions (see
         https://github.com/python-poetry/poetry/pull/8255#issuecomment-1657198242
         for more details).
-        In order to provide results that are as deterministic as possible
-        and consistent between `poetry lock` and `poetry update`, the return value
-        of two different dependencies should not be equal if possible.
         """
-
-        class Preference:
-            """
-            Preference is one of the criteria for choosing which dependency to solve
-            first. A higher value means that there are "more options" to satisfy
-            a dependency. A lower value takes precedence.
-            """
-
-            DIRECT_ORIGIN = 0
-            NO_CHOICE = 1
-            USE_LATEST = 2
-            LOCKED = 3
-            DEFAULT = 4
-
-        def _get_min(dependency: Dependency) -> tuple[int, int, bool, int]:
-            """
-            Returns a tuple of:
-            - preference: see Preference class
-            - num_deps_upper_bound: a dependency with an upper bound is more likely to
-                                    cause conflicts -> a package with more dependencies
-                                    with upper bounds should be chosen first
-            - has_deps: a package with dependencies should be chosen first because
-                        a package without dependencies is less likely to cause conflicts
-            - num_packages: see explanation above
-            """
-            # Direct origin dependencies must be handled first: we don't want to resolve
-            # a regular dependency for some package only to find later that we had a
-            # direct-origin dependency.
-            if dependency.is_direct_origin():
-                return Preference.DIRECT_ORIGIN, 0, False, 0
-
-            use_latest = dependency.name in self._provider.use_latest
-            if not use_latest:
-                locked = self._provider.get_locked(dependency)
-                if locked:
-                    return Preference.LOCKED, 0, False, 0
-
-            packages = self._dependency_cache.search_for(
-                dependency, self._solution.decision_level
-            )
-            num_packages = len(packages)
-            if packages:
-                package = packages[0].package
-                if package.is_root():
-                    relevant_dependencies = package.all_requires
-                else:
-                    if not package.is_direct_origin():
-                        # We have to get the package from the pool,
-                        # otherwise `requires` will be empty.
-                        #
-                        # We might need `package.source_reference` as fallback
-                        # for transitive dependencies without a source
-                        # if there is a top-level dependency
-                        # for the same package with an explicit source.
-                        for repo in (dependency.source_name, package.source_reference):
-                            try:
-                                package = self._provider.get_package_from_pool(
-                                    package.pretty_name,
-                                    package.version,
-                                    repository_name=repo,
-                                )
-                            except Exception:
-                                pass
-                            else:
-                                break
-
-                    relevant_dependencies = [
-                        r
-                        for r in package.requires
-                        if not r.in_extras or r.in_extras[0] in dependency.extras
-                    ]
-                has_deps = bool(relevant_dependencies)
-                num_deps_upper_bound = sum(
-                    1 for d in relevant_dependencies if d.constraint.has_upper_bound()
-                )
+        # Direct origin dependencies must be handled first: we don't want to resolve
+        # a regular dependency for some package only to find later that we had a
+        # direct-origin dependency.
+        if dependency.is_direct_origin():
+            return Preference.DIRECT_ORIGIN, 0, False, 0
+
+        use_latest = dependency.name in self._provider.use_latest
+        if not use_latest:
+            locked = self._provider.get_locked(dependency)
+            if locked:
+                return Preference.LOCKED, 0, False, 0
+
+        packages = self._dependency_cache.search_for(
+            dependency, self._solution.decision_level
+        )
+        num_packages = len(packages)
+        if packages:
+            package = packages[0].package
+            if package.is_root():
+                relevant_dependencies = package.all_requires
             else:
-                has_deps = False
-                num_deps_upper_bound = 0
+                if not package.is_direct_origin():
+                    # We have to get the package from the pool,
+                    # otherwise `requires` will be empty.
+                    #
+                    # We might need `package.source_reference` as fallback
+                    # for transitive dependencies without a source
+                    # if there is a top-level dependency
+                    # for the same package with an explicit source.
+                    for repo in (dependency.source_name, package.source_reference):
+                        try:
+                            package = self._provider.get_package_from_pool(
+                                package.pretty_name,
+                                package.version,
+                                repository_name=repo,
+                            )
+                        except Exception:
+                            pass
+                        else:
+                            break
+
+                relevant_dependencies = [
+                    r
+                    for r in package.requires
+                    if not r.in_extras or r.in_extras[0] in dependency.extras
+                ]
+            has_deps = bool(relevant_dependencies)
+            num_deps_upper_bound = sum(
+                1 for d in relevant_dependencies if d.constraint.has_upper_bound()
+            )
+        else:
+            has_deps = False
+            num_deps_upper_bound = 0
 
-            if num_packages < 2:
-                preference = Preference.NO_CHOICE
-            elif use_latest:
-                preference = Preference.USE_LATEST
-            else:
-                preference = Preference.DEFAULT
-            return preference, -num_deps_upper_bound, not has_deps, -num_packages
+        if num_packages < 2:
+            preference = Preference.NO_CHOICE
+        elif use_latest:
+            preference = Preference.USE_LATEST
+        else:
+            preference = Preference.DEFAULT
+        return preference, -num_deps_upper_bound, not has_deps, -num_packages
 
-        return min(unsatisfied, key=_get_min)
+    def _choose_next(self, unsatisfied: list[Dependency]) -> Dependency:
+        """
+        Chooses the next package to resolve.
+        """
+        return min(unsatisfied, key=self._get_comp_key_cached)
 
     def _choose_package_version(self) -> str | None:
         """