installer: simplify logic for calculating operations regarding extras and fix small bug (#9345)

radoering · radoering · commit 66e124067a73 · 2024-06-02T20:51:22.000+02:00
* Operations are completely calculated in `Transaction.calculate_operations()` so that we do not have to remember extra uninstalls between two calculations.
* Previously, extras that were not requested were not uninstalled when running `poetry install` (without extras) if there was no lockfile, now it does not matter if there was a lockfile or not.
* In `installer._get_extra_packages()` we do not have to distinguish between locked extras and extras in the pyproject.toml because both must be consistent.
diff --git a/src/poetry/installation/installer.py b/src/poetry/installation/installer.py
@@ -7,7 +7,6 @@
 from packaging.utils import canonicalize_name
 
 from poetry.installation.executor import Executor
-from poetry.installation.operations import Install
 from poetry.installation.operations import Uninstall
 from poetry.installation.operations import Update
 from poetry.repositories import Repository
@@ -239,15 +238,17 @@ def _do_install(self) -> int:
                 source_root=self._env.path.joinpath("src")
             ):
                 ops = solver.solve(use_latest=self._whitelist).calculate_operations()
+
+            lockfile_repo = LockfileRepository()
+            self._populate_lockfile_repo(lockfile_repo, ops)
+
         else:
             self._io.write_line("<info>Installing dependencies from lock file</>")
 
-            locked_repository = self._locker.locked_repository()
-
             if not self._locker.is_fresh():
                 raise ValueError(
-                    "pyproject.toml changed significantly since poetry.lock was last generated. "
-                    "Run `poetry lock [--no-update]` to fix the lock file."
+                    "pyproject.toml changed significantly since poetry.lock was last"
+                    " generated. Run `poetry lock [--no-update]` to fix the lock file."
                 )
 
             locker_extras = {
@@ -258,30 +259,25 @@ def _do_install(self) -> int:
                 if extra not in locker_extras:
                     raise ValueError(f"Extra [{extra}] is not specified.")
 
-            # If we are installing from lock
-            # Filter the operations by comparing it with what is
-            # currently installed
-            ops = self._get_operations_from_lock(locked_repository)
-
-        lockfile_repo = LockfileRepository()
-        uninstalls = self._populate_lockfile_repo(lockfile_repo, ops)
+            locked_repository = self._locker.locked_repository()
+            lockfile_repo = locked_repository
 
         if not self.executor.enabled:
             # If we are only in lock mode, no need to go any further
             self._write_lock_file(lockfile_repo)
             return 0
 
-        if self._groups is not None:
-            root = self._package.with_dependency_groups(list(self._groups), only=True)
-        else:
-            root = self._package.without_optional_dependency_groups()
-
         if self._io.is_verbose():
             self._io.write_line("")
             self._io.write_line(
                 "<info>Finding the necessary packages for the current system</>"
             )
 
+        if self._groups is not None:
+            root = self._package.with_dependency_groups(list(self._groups), only=True)
+        else:
+            root = self._package.without_optional_dependency_groups()
+
         # We resolve again by only using the lock file
         packages = lockfile_repo.packages + locked_repository.packages
         pool = RepositoryPool.from_packages(packages, self._config)
@@ -299,31 +295,12 @@ def _do_install(self) -> int:
 
         with solver.use_environment(self._env):
             ops = solver.solve(use_latest=self._whitelist).calculate_operations(
-                with_uninstalls=self._requires_synchronization,
+                with_uninstalls=self._requires_synchronization or self._update,
                 synchronize=self._requires_synchronization,
                 skip_directory=self._skip_directory,
+                extras=set(self._extras),
             )
 
-        if not self._requires_synchronization:
-            # If no packages synchronisation has been requested we need
-            # to calculate the uninstall operations
-            from poetry.puzzle.transaction import Transaction
-
-            transaction = Transaction(
-                locked_repository.packages,
-                [(package, 0) for package in lockfile_repo.packages],
-                installed_packages=self._installed_repository.packages,
-                root_package=root,
-            )
-
-            ops = [
-                op
-                for op in transaction.calculate_operations(with_uninstalls=True)
-                if op.job_type == "uninstall"
-            ] + ops
-        else:
-            ops = uninstalls + ops
-
         # We need to filter operations so that packages
         # not compatible with the current system,
         # or optional and not requested, are dropped
@@ -358,50 +335,15 @@ def _execute(self, operations: list[Operation]) -> int:
 
     def _populate_lockfile_repo(
         self, repo: LockfileRepository, ops: Iterable[Operation]
-    ) -> list[Uninstall]:
-        uninstalls = []
+    ) -> None:
         for op in ops:
             if isinstance(op, Uninstall):
-                uninstalls.append(op)
                 continue
 
             package = op.target_package if isinstance(op, Update) else op.package
             if not repo.has_package(package):
                 repo.add_package(package)
 
-        return uninstalls
-
-    def _get_operations_from_lock(
-        self, locked_repository: Repository
-    ) -> list[Operation]:
-        installed_repo = self._installed_repository
-        ops: list[Operation] = []
-
-        extra_packages = self._get_extra_packages(locked_repository)
-        for locked in locked_repository.packages:
-            is_installed = False
-            for installed in installed_repo.packages:
-                if locked.name == installed.name:
-                    is_installed = True
-                    if locked.optional and locked.name not in extra_packages:
-                        # Installed but optional and not requested in extras
-                        ops.append(Uninstall(locked))
-                    elif locked.version != installed.version:
-                        ops.append(Update(installed, locked))
-
-            # If it's optional and not in required extras
-            # we do not install
-            if locked.optional and locked.name not in extra_packages:
-                continue
-
-            op = Install(locked)
-            if is_installed:
-                op.skip("Already installed")
-
-            ops.append(op)
-
-        return ops
-
     def _filter_operations(self, ops: Iterable[Operation], repo: Repository) -> None:
         extra_packages = self._get_extra_packages(repo)
         for op in ops:
@@ -425,19 +367,11 @@ def _get_extra_packages(self, repo: Repository) -> set[NormalizedName]:
 
         Maybe we just let the solver handle it?
         """
-        extras: dict[NormalizedName, list[NormalizedName]]
-        if self._update:
-            extras = {k: [d.name for d in v] for k, v in self._package.extras.items()}
-        else:
-            raw_extras = self._locker.lock_data.get("extras", {})
-            extras = {
-                canonicalize_name(extra): [
-                    canonicalize_name(dependency) for dependency in dependencies
-                ]
-                for extra, dependencies in raw_extras.items()
-            }
-
-        return get_extra_package_names(repo.packages, extras, self._extras)
+        return get_extra_package_names(
+            repo.packages,
+            {k: [d.name for d in v] for k, v in self._package.extras.items()},
+            self._extras,
+        )
 
     def _get_installed(self) -> InstalledRepository:
         return InstalledRepository.load(self._env)
diff --git a/src/poetry/puzzle/transaction.py b/src/poetry/puzzle/transaction.py
@@ -2,8 +2,11 @@
 
 from typing import TYPE_CHECKING
 
+from poetry.utils.extras import get_extra_package_names
+
 
 if TYPE_CHECKING:
+    from packaging.utils import NormalizedName
     from poetry.core.packages.package import Package
 
     from poetry.installation.operations.operation import Operation
@@ -32,23 +35,42 @@ def calculate_operations(
         synchronize: bool = False,
         *,
         skip_directory: bool = False,
+        extras: set[NormalizedName] | None = None,
     ) -> list[Operation]:
         from poetry.installation.operations import Install
         from poetry.installation.operations import Uninstall
         from poetry.installation.operations import Update
 
         operations: list[Operation] = []
 
+        extra_packages: set[NormalizedName] = set()
+        if extras is not None:
+            assert self._root_package is not None
+            extra_packages = get_extra_package_names(
+                [package for package, _ in self._result_packages],
+                {k: [d.name for d in v] for k, v in self._root_package.extras.items()},
+                extras,
+            )
+
+        uninstalls: set[NormalizedName] = set()
         for result_package, priority in self._result_packages:
             installed = False
 
             for installed_package in self._installed_packages:
                 if result_package.name == installed_package.name:
                     installed = True
 
+                    # Extras that were not requested are always uninstalled.
+                    if extras is not None and (
+                        result_package.optional
+                        and result_package.name not in extra_packages
+                    ):
+                        uninstalls.add(installed_package.name)
+                        operations.append(Uninstall(installed_package))
+
                     # We have to perform an update if the version or another
                     # attribute of the package has changed (source type, url, ref, ...).
-                    if result_package.version != installed_package.version or (
+                    elif result_package.version != installed_package.version or (
                         (
                             # This has to be done because installed packages cannot
                             # have type "legacy". If a package with type "legacy"
@@ -81,7 +103,6 @@ def calculate_operations(
                 operations.append(Install(result_package, priority=priority))
 
         if with_uninstalls:
-            uninstalls: set[str] = set()
             for current_package in self._current_packages:
                 found = any(
                     current_package.name == result_package.name
@@ -92,7 +113,7 @@ def calculate_operations(
                     for installed_package in self._installed_packages:
                         if installed_package.name == current_package.name:
                             uninstalls.add(installed_package.name)
-                            operations.append(Uninstall(current_package))
+                            operations.append(Uninstall(installed_package))
 
             if synchronize:
                 result_package_names = {
diff --git a/tests/installation/test_installer.py b/tests/installation/test_installer.py
@@ -1089,7 +1089,7 @@ def test_run_installs_extras_with_deps_if_requested(
         expected_installations_count = 0 if is_installed else 2
         # We only want to uninstall extras if we do a "poetry install" without extras,
         # not if we do a "poetry update" or "poetry add".
-        expected_removals_count = 2 if is_installed and is_locked else 0
+        expected_removals_count = 2 if is_installed else 0
 
     assert installer.executor.installations_count == expected_installations_count
     assert installer.executor.removals_count == expected_removals_count
