From 4ed5e1a84aa7c1f3431b5ac35ccab81dbd3e9807 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Randy=20D=C3=B6ring?=
 <30527984+radoering@users.noreply.github.com>
Date: Tue, 13 Feb 2024 16:55:44 +0100
Subject: [PATCH] refactor: use Package instead of DependencyPackage as key in
 overrides dict (#8956)

---
 src/poetry/puzzle/exceptions.py |  8 +++-----
 src/poetry/puzzle/provider.py   | 22 +++++++---------------
 src/poetry/puzzle/solver.py     |  5 ++---
 3 files changed, 12 insertions(+), 23 deletions(-)

diff --git a/src/poetry/puzzle/exceptions.py b/src/poetry/puzzle/exceptions.py
index 58087c4a78d..6bb7c0027bc 100644
--- a/src/poetry/puzzle/exceptions.py
+++ b/src/poetry/puzzle/exceptions.py
@@ -5,9 +5,9 @@
 
 if TYPE_CHECKING:
     from poetry.core.packages.dependency import Dependency
+    from poetry.core.packages.package import Package
 
     from poetry.mixology.failure import SolveFailure
-    from poetry.packages import DependencyPackage
 
 
 class SolverProblemError(Exception):
@@ -22,11 +22,9 @@ def error(self) -> SolveFailure:
 
 
 class OverrideNeeded(Exception):
-    def __init__(
-        self, *overrides: dict[DependencyPackage, dict[str, Dependency]]
-    ) -> None:
+    def __init__(self, *overrides: dict[Package, dict[str, Dependency]]) -> None:
         self._overrides = overrides
 
     @property
-    def overrides(self) -> tuple[dict[DependencyPackage, dict[str, Dependency]], ...]:
+    def overrides(self) -> tuple[dict[Package, dict[str, Dependency]], ...]:
         return self._overrides
diff --git a/src/poetry/puzzle/provider.py b/src/poetry/puzzle/provider.py
index dc39ee11469..d404f84d520 100644
--- a/src/poetry/puzzle/provider.py
+++ b/src/poetry/puzzle/provider.py
@@ -125,7 +125,7 @@ def __init__(
         self._env: Env | None = None
         self._python_constraint = package.python_constraint
         self._is_debugging: bool = self._io.is_debug() or self._io.is_very_verbose()
-        self._overrides: dict[DependencyPackage, dict[str, Dependency]] = {}
+        self._overrides: dict[Package, dict[str, Dependency]] = {}
         self._deferred_cache: dict[Dependency, Package] = {}
         self._load_deferred = True
         self._source_root: Path | None = None
@@ -155,9 +155,7 @@ def use_latest(self) -> Collection[NormalizedName]:
     def is_debugging(self) -> bool:
         return self._is_debugging
 
-    def set_overrides(
-        self, overrides: dict[DependencyPackage, dict[str, Dependency]]
-    ) -> None:
+    def set_overrides(self, overrides: dict[Package, dict[str, Dependency]]) -> None:
         self._overrides = overrides
 
     def load_deferred(self, load_deferred: bool) -> None:
@@ -383,7 +381,7 @@ def _search_for_url(self, dependency: URLDependency) -> Package:
         return package
 
     def _get_dependencies_with_overrides(
-        self, dependencies: list[Dependency], package: DependencyPackage
+        self, dependencies: list[Dependency], package: Package
     ) -> list[Dependency]:
         overrides = self._overrides.get(package, {})
         _dependencies = []
@@ -458,9 +456,7 @@ def incompatibilities_for(
             and self._python_constraint.allows_any(dep.python_constraint)
             and (not self._env or dep.marker.validate(self._env.marker_env))
         ]
-        dependencies = self._get_dependencies_with_overrides(
-            _dependencies, dependency_package
-        )
+        dependencies = self._get_dependencies_with_overrides(_dependencies, package)
 
         return [
             Incompatibility(
@@ -568,9 +564,7 @@ def complete_package(
                         continue
                     self.search_for_direct_origin_dependency(dep)
 
-        dependencies = self._get_dependencies_with_overrides(
-            _dependencies, dependency_package
-        )
+        dependencies = self._get_dependencies_with_overrides(_dependencies, package)
 
         # Searching for duplicate dependencies
         #
@@ -647,11 +641,9 @@ def fmt_warning(d: Dependency) -> str:
             for dep in deps:
                 if not overrides_marker_intersection.intersect(dep.marker).is_empty():
                     current_overrides = self._overrides.copy()
-                    package_overrides = current_overrides.get(
-                        dependency_package, {}
-                    ).copy()
+                    package_overrides = current_overrides.get(package, {}).copy()
                     package_overrides.update({dep.name: dep})
-                    current_overrides.update({dependency_package: package_overrides})
+                    current_overrides.update({package: package_overrides})
                     overrides.append(current_overrides)
 
             if overrides:
diff --git a/src/poetry/puzzle/solver.py b/src/poetry/puzzle/solver.py
index 3e0b2d4d363..af107646876 100644
--- a/src/poetry/puzzle/solver.py
+++ b/src/poetry/puzzle/solver.py
@@ -27,7 +27,6 @@
     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 RepositoryPool
     from poetry.utils.env import Env
@@ -51,7 +50,7 @@ def __init__(
         self._provider = Provider(
             self._package, self._pool, self._io, installed=installed, locked=locked
         )
-        self._overrides: list[dict[DependencyPackage, dict[str, Dependency]]] = []
+        self._overrides: list[dict[Package, dict[str, Dependency]]] = []
 
     @property
     def provider(self) -> Provider:
@@ -119,7 +118,7 @@ def _progress(self) -> Iterator[None]:
 
     def _solve_in_compatibility_mode(
         self,
-        overrides: tuple[dict[DependencyPackage, dict[str, Dependency]], ...],
+        overrides: tuple[dict[Package, dict[str, Dependency]], ...],
     ) -> tuple[list[Package], list[int]]:
         packages = []
         depths = []