From ebf88523cbcbfe7afcd74bed6a705c898c80b181 Mon Sep 17 00:00:00 2001
From: lucemia <lucemia@gmail.com>
Date: Mon, 18 Dec 2023 15:39:49 +0800
Subject: [PATCH] fix #8614 source constriants failed with extras

---
 src/poetry/poetry.py                       | 13 ++++++++++++-
 src/poetry/repositories/repository_pool.py | 10 ++++++++++
 tests/repositories/test_repository_pool.py | 17 +++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/poetry/poetry.py b/src/poetry/poetry.py
index dfbe043594c..f5b070e95ce 100644
--- a/src/poetry/poetry.py
+++ b/src/poetry/poetry.py
@@ -41,7 +41,18 @@ def __init__(
 
         self._locker = locker
         self._config = config
-        self._pool = RepositoryPool(config=config)
+
+        local_config = local_config or {}
+        dependency_source_cache = {}
+
+        for group in [*local_config.get("group", {}).values(), local_config]:
+            for name, dependency in group.get("dependencies", {}).items():
+                if isinstance(dependency, dict) and "source" in dependency:
+                    dependency_source_cache[name] = dependency["source"]
+
+        self._pool = RepositoryPool(
+            config=config, dependency_source_mapping=dependency_source_cache
+        )
         self._plugin_manager: PluginManager | None = None
         self._disable_cache = disable_cache
 
diff --git a/src/poetry/repositories/repository_pool.py b/src/poetry/repositories/repository_pool.py
index 9d6338f4ff2..3d8d77b3638 100644
--- a/src/poetry/repositories/repository_pool.py
+++ b/src/poetry/repositories/repository_pool.py
@@ -45,6 +45,7 @@ def __init__(
         repositories: list[Repository] | None = None,
         ignore_repository_names: object = _SENTINEL,
         *,
+        dependency_source_mapping: dict[str, str] | None = None,
         config: Config | None = None,
     ) -> None:
         super().__init__("poetry-repository-pool")
@@ -59,6 +60,10 @@ def __init__(
             cache_dir=(config or Config.create()).artifacts_cache_directory
         )
 
+        self._dependency_source_mapping: dict[str, str] = (
+            dependency_source_mapping or {}
+        )
+
         if ignore_repository_names is not _SENTINEL:
             warnings.warn(
                 "The 'ignore_repository_names' argument to 'RepositoryPool.__init__' is"
@@ -211,6 +216,11 @@ def find_packages(self, dependency: Dependency) -> list[Package]:
         if repository_name:
             return self.repository(repository_name).find_packages(dependency)
 
+        if dependency.name in self._dependency_source_mapping:
+            return self.repository(
+                self._dependency_source_mapping[dependency.name]
+            ).find_packages(dependency)
+
         packages: list[Package] = []
         for repo in self.repositories:
             if packages and self.get_priority(repo.name) is Priority.SUPPLEMENTAL:
diff --git a/tests/repositories/test_repository_pool.py b/tests/repositories/test_repository_pool.py
index 2a62f360815..a2bdee681ac 100644
--- a/tests/repositories/test_repository_pool.py
+++ b/tests/repositories/test_repository_pool.py
@@ -333,6 +333,23 @@ def test_pool_find_packages_in_specified_repository() -> None:
     assert returned_packages_unavailable == []
 
 
+def test_pool_find_packages_with_dependency_source_mapping() -> None:
+    package_foo1 = get_package("foo1", "1.1.1")
+    package_foo2 = get_package("foo2", "1.1.1")
+
+    bar = Repository("bar", [package_foo1, package_foo2])
+    pool = RepositoryPool(dependency_source_mapping={"foo1": "bar"})
+    pool.add_repository(bar, priority=Priority.EXPLICIT)
+
+    available_dependency = get_dependency("foo1", "^1.0.0")
+    returned_packages_available = pool.find_packages(available_dependency)
+    unavailable_dependency = get_dependency("foo2", "^1.0.0")
+    returned_unavailable_dependency = pool.find_packages(unavailable_dependency)
+
+    assert returned_packages_available == [package_foo1]
+    assert returned_unavailable_dependency == []
+
+
 def test_search_no_legacy_repositories() -> None:
     package_foo1 = get_package("foo", "1.0.0")
     package_foo2 = get_package("foo", "2.0.0")