refactor(solver/provider/version_solver): move get_locked() from Vers…

…ionSolver to Provider

src/poetry/installation/installer.py

@@ -23,6 +23,7 @@
     from collections.abc import Sequence
 
     from cleo.io.io import IO
+    from packaging.utils import NormalizedName
     from poetry.core.packages.project_package import ProjectPackage
 
     from poetry.config.config import Config
@@ -60,7 +61,7 @@ def __init__(
         self._execute_operations = True
         self._lock = False
 
-        self._whitelist: list[str] = []
+        self._whitelist: list[NormalizedName] = []
 
         self._extras: list[str] = []
 

src/poetry/mixology/__init__.py

@@ -9,16 +9,10 @@
     from poetry.core.packages.project_package import ProjectPackage
 
     from poetry.mixology.result import SolverResult
-    from poetry.packages import DependencyPackage
     from poetry.puzzle.provider import Provider
 
 
-def resolve_version(
-    root: ProjectPackage,
-    provider: Provider,
-    locked: dict[str, list[DependencyPackage]] | None = None,
-    use_latest: list[str] | None = None,
-) -> SolverResult:
-    solver = VersionSolver(root, provider, locked=locked, use_latest=use_latest)
+def resolve_version(root: ProjectPackage, provider: Provider) -> SolverResult:
+    solver = VersionSolver(root, provider)
 
     return solver.solve()

src/poetry/mixology/version_solver.py

@@ -18,12 +18,12 @@
 from poetry.mixology.result import SolverResult
 from poetry.mixology.set_relation import SetRelation
 from poetry.mixology.term import Term
-from poetry.packages import DependencyPackage
 
 
 if TYPE_CHECKING:
     from poetry.core.packages.project_package import ProjectPackage
 
+    from poetry.packages import DependencyPackage
     from poetry.puzzle.provider import Provider
 
 
@@ -82,23 +82,10 @@ class VersionSolver:
     on how this solver works.
     """
 
-    def __init__(
-        self,
-        root: ProjectPackage,
-        provider: Provider,
-        locked: dict[str, list[DependencyPackage]] | None = None,
-        use_latest: list[str] | None = None,
-    ) -> None:
+    def __init__(self, root: ProjectPackage, provider: Provider) -> None:
         self._root = root
         self._provider = provider
         self._dependency_cache = DependencyCache(provider)
-        self._locked = locked or {}
-
-        if use_latest is None:
-            use_latest = []
-
-        self._use_latest = use_latest
-
         self._incompatibilities: dict[str, list[Incompatibility]] = {}
         self._contradicted_incompatibilities: set[Incompatibility] = set()
         self._solution = PartialSolution()
@@ -384,12 +371,12 @@ def _get_min(dependency: Dependency) -> tuple[bool, int]:
             if dependency.is_direct_origin():
                 return False, -1
 
-            if dependency.name in self._use_latest:
+            if dependency.name in self._provider.use_latest:
                 # If we're forced to use the latest version of a package, it effectively
                 # only has one version to choose from.
                 return not dependency.marker.is_any(), 1
 
-            locked = self._get_locked(dependency)
+            locked = self._provider.get_locked(dependency)
             if locked:
                 return not dependency.marker.is_any(), 1
 
@@ -406,7 +393,7 @@ def _get_min(dependency: Dependency) -> tuple[bool, int]:
         else:
             dependency = min(*unsatisfied, key=_get_min)
 
-        locked = self._get_locked(dependency)
+        locked = self._provider.get_locked(dependency)
         if locked is None:
             try:
                 packages = self._dependency_cache.search_for(dependency)
@@ -499,23 +486,5 @@ def _add_incompatibility(self, incompatibility: Incompatibility) -> None:
                 incompatibility
             )
 
-    def _get_locked(self, dependency: Dependency) -> DependencyPackage | None:
-        if dependency.name in self._use_latest:
-            return None
-
-        locked = self._locked.get(dependency.name, [])
-        for dependency_package in locked:
-            package = dependency_package.package
-            if (
-                # Locked dependencies are always without features.
-                # Thus, we can't use is_same_package_as() here because it compares
-                # the complete_name (including features).
-                dependency.name == package.name
-                and dependency.is_same_source_as(package)
-                and dependency.constraint.allows(package.version)
-            ):
-                return DependencyPackage(dependency, package)
-        return None
-
     def _log(self, text: str) -> None:
         self._provider.debug(text, self._solution.attempted_solutions)

src/poetry/puzzle/provider.py

@@ -12,6 +12,7 @@
 from contextlib import contextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING
+from typing import Collection
 from typing import cast
 
 from cleo.ui.progress_indicator import ProgressIndicator
@@ -42,6 +43,7 @@
     from collections.abc import Iterator
 
     from cleo.io.io import IO
+    from packaging.utils import NormalizedName
     from poetry.core.packages.dependency import Dependency
     from poetry.core.packages.directory_dependency import DirectoryDependency
     from poetry.core.packages.file_dependency import FileDependency
@@ -127,6 +129,7 @@ def __init__(
         io: IO,
         *,
         installed: list[Package] | None = None,
+        locked: list[Package] | None = None,
     ) -> None:
         self._package = package
         self._pool = pool
@@ -140,11 +143,27 @@ def __init__(
         self._source_root: Path | None = None
         self._installed_packages = installed if installed is not None else []
         self._direct_origin_packages: dict[str, Package] = {}
+        self._locked: dict[NormalizedName, list[DependencyPackage]] = defaultdict(list)
+        self._use_latest: Collection[NormalizedName] = []
+
+        for package in locked or []:
+            self._locked[package.name].append(
+                DependencyPackage(package.to_dependency(), package)
+            )
+        for dependency_packages in self._locked.values():
+            dependency_packages.sort(
+                key=lambda p: p.package.version,
+                reverse=True,
+            )
 
     @property
     def pool(self) -> Pool:
         return self._pool
 
+    @property
+    def use_latest(self) -> Collection[NormalizedName]:
+        return self._use_latest
+
     def is_debugging(self) -> bool:
         return self._is_debugging
 
@@ -161,9 +180,10 @@ def use_source_root(self, source_root: Path) -> Iterator[Provider]:
         original_source_root = self._source_root
         self._source_root = source_root
 
-        yield self
-
-        self._source_root = original_source_root
+        try:
+            yield self
+        finally:
+            self._source_root = original_source_root
 
     @contextmanager
     def use_environment(self, env: Env) -> Iterator[Provider]:
@@ -172,10 +192,20 @@ def use_environment(self, env: Env) -> Iterator[Provider]:
         self._env = env
         self._python_constraint = Version.parse(env.marker_env["python_full_version"])
 
-        yield self
+        try:
+            yield self
+        finally:
+            self._env = None
+            self._python_constraint = original_python_constraint
 
-        self._env = None
-        self._python_constraint = original_python_constraint
+    @contextmanager
+    def use_latest_for(self, names: Collection[NormalizedName]) -> Iterator[Provider]:
+        self._use_latest = names
+
+        try:
+            yield self
+        finally:
+            self._use_latest = []
 
     @staticmethod
     def validate_package_for_dependency(
@@ -801,6 +831,24 @@ def fmt_warning(d: Dependency) -> str:
 
         return dependency_package
 
+    def get_locked(self, dependency: Dependency) -> DependencyPackage | None:
+        if dependency.name in self._use_latest:
+            return None
+
+        locked = self._locked.get(dependency.name, [])
+        for dependency_package in locked:
+            package = dependency_package.package
+            if (
+                # Locked dependencies are always without features.
+                # Thus, we can't use is_same_package_as() here because it compares
+                # the complete_name (including features).
+                dependency.name == package.name
+                and dependency.is_same_source_as(package)
+                and dependency.constraint.allows(package.version)
+            ):
+                return DependencyPackage(dependency, package)
+        return None
+
     def debug(self, message: str, depth: int = 0) -> None:
         if not (self._io.is_very_verbose() or self._io.is_debug()):
             return

src/poetry/puzzle/solver.py

@@ -5,6 +5,7 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from typing import TYPE_CHECKING
+from typing import Collection
 from typing import FrozenSet
 from typing import Tuple
 from typing import TypeVar
@@ -13,7 +14,6 @@
 
 from poetry.mixology import resolve_version
 from poetry.mixology.failure import SolveFailure
-from poetry.packages import DependencyPackage
 from poetry.puzzle.exceptions import OverrideNeeded
 from poetry.puzzle.exceptions import SolverProblemError
 from poetry.puzzle.provider import Indicator
@@ -24,10 +24,12 @@
     from collections.abc import Iterator
 
     from cleo.io.io import IO
+    from packaging.utils import NormalizedName
     from poetry.core.packages.dependency import Dependency
     from poetry.core.packages.package import Package
     from poetry.core.packages.project_package import ProjectPackage
 
+    from poetry.packages import DependencyPackage
     from poetry.puzzle.transaction import Transaction
     from poetry.repositories import Pool
     from poetry.utils.env import Env
@@ -49,7 +51,7 @@ def __init__(
         self._io = io
 
         self._provider = Provider(
-            self._package, self._pool, self._io, installed=installed
+            self._package, self._pool, self._io, installed=installed, locked=locked
         )
         self._overrides: list[dict[DependencyPackage, dict[str, Dependency]]] = []
 
@@ -62,12 +64,14 @@ def use_environment(self, env: Env) -> Iterator[None]:
         with self.provider.use_environment(env):
             yield
 
-    def solve(self, use_latest: list[str] | None = None) -> Transaction:
+    def solve(
+        self, use_latest: Collection[NormalizedName] | None = None
+    ) -> Transaction:
         from poetry.puzzle.transaction import Transaction
 
-        with self._progress():
+        with self._progress(), self._provider.use_latest_for(use_latest or []):
             start = time.time()
-            packages, depths = self._solve(use_latest=use_latest)
+            packages, depths = self._solve()
             end = time.time()
 
             if len(self._overrides) > 1:
@@ -116,7 +120,6 @@ def _progress(self) -> Iterator[None]:
     def _solve_in_compatibility_mode(
         self,
         overrides: tuple[dict[DependencyPackage, dict[str, Dependency]], ...],
-        use_latest: list[str] | None = None,
     ) -> tuple[list[Package], list[int]]:
         packages = []
         depths = []
@@ -126,7 +129,7 @@ def _solve_in_compatibility_mode(
                 f"with the following overrides ({override}).</comment>"
             )
             self._provider.set_overrides(override)
-            _packages, _depths = self._solve(use_latest=use_latest)
+            _packages, _depths = self._solve()
             for index, package in enumerate(_packages):
                 if package not in packages:
                     packages.append(package)
@@ -143,31 +146,16 @@ def _solve_in_compatibility_mode(
 
         return packages, depths
 
-    def _solve(
-        self, use_latest: list[str] | None = None
-    ) -> tuple[list[Package], list[int]]:
+    def _solve(self) -> tuple[list[Package], list[int]]:
         if self._provider._overrides:
             self._overrides.append(self._provider._overrides)
 
-        locked: dict[str, list[DependencyPackage]] = defaultdict(list)
-        for package in self._locked_packages:
-            locked[package.name].append(
-                DependencyPackage(package.to_dependency(), package)
-            )
-        for dependency_packages in locked.values():
-            dependency_packages.sort(
-                key=lambda p: p.package.version,
-                reverse=True,
-            )
-
         try:
-            result = resolve_version(
-                self._package, self._provider, locked=locked, use_latest=use_latest
-            )
+            result = resolve_version(self._package, self._provider)
 
             packages = result.packages
         except OverrideNeeded as e:
-            return self._solve_in_compatibility_mode(e.overrides, use_latest=use_latest)
+            return self._solve_in_compatibility_mode(e.overrides)
         except SolveFailure as e:
             raise SolverProblemError(e)