refactor: simplify Pool logic

Introduces Priority enum to help in bookkeeping of repository types.
Additionally specifies parameter 'repository' to 'repository_name' where
relevant.

Author: b-kamphorst

src/poetry/puzzle/provider.py

@@ -555,7 +555,7 @@ def complete_package(
                         package.pretty_name,
                         package.version,
                         extras=list(dependency.extras),
-                        repository=dependency.source_name,
+                        repository_name=dependency.source_name,
                     ),
                 )
             except PackageNotFound as e:

src/poetry/repositories/legacy_repository.py

@@ -33,6 +33,18 @@ def __init__(
 
         super().__init__(name, url.rstrip("/"), config, disable_cache)
 
+    @property
+    def packages(self) -> list[Package]:
+        # LegacyRepository._packages is not populated and other implementations
+        # implicitly rely on this (e.g. Pool.search via
+        # LegacyRepository.search). To avoid special-casing Pool or changing
+        # behavior, we stub and return an empty list.
+        #
+        # TODO: Rethinking search behaviour and design.
+        # Ref: https://github.com/python-poetry/poetry/issues/2446 and
+        # https://github.com/python-poetry/poetry/pull/6669#discussion_r990874908.
+        return []
+
     def package(
         self, name: str, version: Version, extras: list[str] | None = None
     ) -> Package:

src/poetry/repositories/pool.py

@@ -1,5 +1,9 @@
 from __future__ import annotations
 
+import enum
+
+from collections import OrderedDict
+from enum import IntEnum
 from typing import TYPE_CHECKING
 
 from poetry.repositories.exceptions import PackageNotFound
@@ -13,53 +17,60 @@
     from poetry.repositories.repository import Repository
 
 
+class Priority(IntEnum):
+    # The order of the members below dictates the actual priority. The first member has
+    # top priority.
+    DEFAULT = enum.auto()
+    PRIMARY = enum.auto()
+    SECONDARY = enum.auto()
+
+
 class Pool:
     def __init__(
         self,
         repositories: list[Repository] | None = None,
         ignore_repository_names: bool = False,
     ) -> None:
         self._name = "poetry-pool"
+        self._repositories: OrderedDict[
+            str, tuple[Repository, RepositoryPriority]
+        ] = OrderedDict()
+        self._ignore_repository_names = ignore_repository_names
 
         if repositories is None:
             repositories = []
-
-        self._lookup: dict[str, int] = {}
-        self._repositories: list[Repository] = []
-        self._default = False
-        self._has_primary_repositories = False
-        self._secondary_start_idx: int | None = None
-
         for repository in repositories:
             self.add_repository(repository)
 
-        self._ignore_repository_names = ignore_repository_names
-
     @property
     def name(self) -> str:
         return self._name
 
     @property
     def repositories(self) -> list[Repository]:
-        return self._repositories
+        unsorted_repositories = self._repositories.values()
+        sorted_repositories = sorted(unsorted_repositories, key=lambda p: p[1].value)
+        return [repo for (repo, _) in sorted_repositories]
 
     def has_default(self) -> bool:
-        return self._default
+        return self._contains_priority(Priority.DEFAULT)
 
     def has_primary_repositories(self) -> bool:
-        return self._has_primary_repositories
+        return self._contains_priority(Priority.PRIMARY)
+
+    def _contains_priority(self, priority: Priority) -> bool:
+        return any(
+            repo_prio is priority for _, repo_prio in self._repositories.values()
+        )
 
     def has_repository(self, name: str) -> bool:
-        return name.lower() in self._lookup
+        return name.lower() in self._repositories
 
     def repository(self, name: str) -> Repository:
         name = name.lower()
-
-        lookup = self._lookup.get(name)
-        if lookup is not None:
-            return self._repositories[lookup]
-
-        raise ValueError(f'Repository "{name}" does not exist.')
+        if self.has_repository(name):
+            return self._repositories[name][0]
+        raise IndexError(f'Repository "{name}" does not exist.')
 
     def add_repository(
         self, repository: Repository, default: bool = False, secondary: bool = False
@@ -68,130 +79,59 @@ def add_repository(
         Adds a repository to the pool.
         """
         repository_name = repository.name.lower()
-        if repository_name in self._lookup:
-            raise ValueError(f"{repository_name} already added")
-
-        if default:
-            if self.has_default():
-                raise ValueError("Only one repository can be the default")
-
-            self._default = True
-            self._repositories.insert(0, repository)
-            for name in self._lookup:
-                self._lookup[name] += 1
+        if self.has_repository(repository_name):
+            raise ValueError(
+                f"A repository with name {repository_name} was already added."
+            )
 
-            if self._secondary_start_idx is not None:
-                self._secondary_start_idx += 1
+        if default and self.has_default():
+            raise ValueError("Only one repository can be the default.")
 
-            self._lookup[repository_name] = 0
+        priority = Priority.PRIMARY
+        if default:
+            priority = Priority.DEFAULT
         elif secondary:
-            if self._secondary_start_idx is None:
-                self._secondary_start_idx = len(self._repositories)
-
-            self._repositories.append(repository)
-            self._lookup[repository_name] = len(self._repositories) - 1
-        else:
-            self._has_primary_repositories = True
-            if self._secondary_start_idx is None:
-                self._repositories.append(repository)
-                self._lookup[repository_name] = len(self._repositories) - 1
-            else:
-                self._repositories.insert(self._secondary_start_idx, repository)
-
-                for name, idx in self._lookup.items():
-                    if idx < self._secondary_start_idx:
-                        continue
-
-                    self._lookup[name] += 1
-
-                self._lookup[repository_name] = self._secondary_start_idx
-                self._secondary_start_idx += 1
-
+            priority = Priority.SECONDARY
+        self._repositories[repository_name] = (repository, priority)
         return self
 
-    def remove_repository(self, repository_name: str) -> Pool:
-        if repository_name is not None:
-            repository_name = repository_name.lower()
-
-        idx = self._lookup.get(repository_name)
-        if idx is not None:
-            del self._repositories[idx]
-            del self._lookup[repository_name]
-
-            if idx == 0:
-                self._default = False
-
-            for name in self._lookup:
-                if self._lookup[name] > idx:
-                    self._lookup[name] -= 1
-
-            if (
-                self._secondary_start_idx is not None
-                and self._secondary_start_idx > idx
-            ):
-                self._secondary_start_idx -= 1
-
+    def remove_repository(self, name: str) -> Pool:
+        if not self.has_repository(name):
+            raise IndexError(f"Pool can not remove unknown repository '{name}'.")
+        del self._repositories[name.lower()]
         return self
 
     def package(
         self,
         name: str,
         version: Version,
         extras: list[str] | None = None,
-        repository: str | None = None,
+        repository_name: str | None = None,
     ) -> Package:
-        if repository is not None:
-            repository = repository.lower()
-
-        if (
-            repository is not None
-            and repository not in self._lookup
-            and not self._ignore_repository_names
-        ):
-            raise ValueError(f'Repository "{repository}" does not exist.')
-
-        if repository is not None and not self._ignore_repository_names:
-            return self.repository(repository).package(name, version, extras=extras)
+        if repository_name and not self._ignore_repository_names:
+            return self.repository(repository_name).package(
+                name, version, extras=extras
+            )
 
-        for repo in self._repositories:
+        for repo in self.repositories:
             try:
-                package = repo.package(name, version, extras=extras)
+                return repo.package(name, version, extras=extras)
             except PackageNotFound:
                 continue
-
-            return package
-
         raise PackageNotFound(f"Package {name} ({version}) not found.")
 
     def find_packages(self, dependency: Dependency) -> list[Package]:
-        repository = dependency.source_name
-        if repository is not None:
-            repository = repository.lower()
-
-        if (
-            repository is not None
-            and repository not in self._lookup
-            and not self._ignore_repository_names
-        ):
-            raise ValueError(f'Repository "{repository}" does not exist.')
-
-        if repository is not None and not self._ignore_repository_names:
-            return self.repository(repository).find_packages(dependency)
-
-        packages = []
-        for repo in self._repositories:
-            packages += repo.find_packages(dependency)
+        repository_name = dependency.source_name
+        if repository_name and not self._ignore_repository_names:
+            return self.repository(repository_name).find_packages(dependency)
 
+        packages: list[Package] = []
+        for repo in self.repositories:
+            packages += repo.find_packages(dependency)
         return packages
 
     def search(self, query: str) -> list[Package]:
-        from poetry.repositories.legacy_repository import LegacyRepository
-
-        results = []
-        for repository in self._repositories:
-            if isinstance(repository, LegacyRepository):
-                continue
-
+        results: list[Package] = []
+        for repository in self.repositories:
             results += repository.search(query)
-
         return results

tests/console/commands/test_add.py

@@ -858,7 +858,7 @@ def test_add_constraint_with_source(
 def test_add_constraint_with_source_that_does_not_exist(
     app: PoetryTestApplication, tester: CommandTester
 ):
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(IndexError) as e:
         tester.execute("foo --source i-dont-exist")
 
     assert str(e.value) == 'Repository "i-dont-exist" does not exist.'
@@ -1848,7 +1848,7 @@ def test_add_constraint_with_source_old_installer(
 def test_add_constraint_with_source_that_does_not_exist_old_installer(
     app: PoetryTestApplication, old_tester: CommandTester
 ):
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(IndexError) as e:
         old_tester.execute("foo --source i-dont-exist")
 
     assert str(e.value) == 'Repository "i-dont-exist" does not exist.'

tests/repositories/test_legacy_repository.py

@@ -63,6 +63,13 @@ def _download(self, url: str, dest: Path) -> None:
         shutil.copyfile(str(filepath), dest)
 
 
+def test_packages_property_returns_empty_list() -> None:
+    repo = MockRepository()
+    repo._packages = [repo.package("jupyter", Version.parse("1.0.0"))]
+
+    assert repo.packages == []
+
+
 def test_page_relative_links_path_are_correct() -> None:
     repo = MockRepository()
 

tests/repositories/test_pool.py

@@ -32,7 +32,7 @@ def test_pool_with_initial_repositories() -> None:
 def test_repository_no_repository() -> None:
     pool = Pool()
 
-    with pytest.raises(ValueError):
+    with pytest.raises(IndexError):
         pool.repository("foo")
 
 
@@ -178,7 +178,9 @@ def test_pool_get_package_in_specified_repository() -> None:
     repo2 = Repository("repo2", [package])
     pool = Pool([repo1, repo2])
 
-    returned_package = pool.package("foo", Version.parse("1.0.0"), repository="repo2")
+    returned_package = pool.package(
+        "foo", Version.parse("1.0.0"), repository_name="repo2"
+    )
 
     assert returned_package == package
 
@@ -198,7 +200,7 @@ def test_pool_no_package_from_specified_repository_raises_package_not_found() ->
     pool = Pool([repo1, repo2])
 
     with pytest.raises(PackageNotFound):
-        pool.package("foo", Version.parse("1.0.0"), repository="repo1")
+        pool.package("foo", Version.parse("1.0.0"), repository_name="repo1")
 
 
 def test_pool_find_packages_in_any_repository() -> None:

tests/test_factory.py

@@ -299,7 +299,7 @@ def test_poetry_with_two_default_sources(with_simple_keyring: None):
     with pytest.raises(ValueError) as e:
         Factory().create_poetry(fixtures_dir / "with_two_default_sources")
 
-    assert str(e.value) == "Only one repository can be the default"
+    assert str(e.value) == "Only one repository can be the default."
 
 
 def test_validate():