locker: lock transitive marker and groups for each package (python-poetry#9427)

Author: radoering

src/poetry/installation/installer.py

@@ -7,8 +7,6 @@
 from packaging.utils import canonicalize_name
 
 from poetry.installation.executor import Executor
-from poetry.installation.operations import Uninstall
-from poetry.installation.operations import Update
 from poetry.repositories import Repository
 from poetry.repositories import RepositoryPool
 from poetry.repositories.installed_repository import InstalledRepository
@@ -20,12 +18,14 @@
 
     from cleo.io.io import IO
     from packaging.utils import NormalizedName
+    from poetry.core.packages.package import Package
     from poetry.core.packages.path_dependency import PathDependency
     from poetry.core.packages.project_package import ProjectPackage
 
     from poetry.config.config import Config
     from poetry.installation.operations.operation import Operation
     from poetry.packages import Locker
+    from poetry.packages.transitive_package_info import TransitivePackageInfo
     from poetry.utils.env import Env
 
 
@@ -196,12 +196,9 @@ def _do_refresh(self) -> int:
         with solver.provider.use_source_root(
             source_root=self._env.path.joinpath("src")
         ):
-            ops = solver.solve(use_latest=use_latest).calculate_operations()
+            solved_packages = solver.solve(use_latest=use_latest).get_solved_packages()
 
-        lockfile_repo = LockfileRepository()
-        self._populate_lockfile_repo(lockfile_repo, ops)
-
-        self._write_lock_file(lockfile_repo, force=True)
+        self._write_lock_file(solved_packages, force=True)
 
         return 0
 
@@ -236,10 +233,18 @@ def _do_install(self) -> int:
             with solver.provider.use_source_root(
                 source_root=self._env.path.joinpath("src")
             ):
-                ops = solver.solve(use_latest=self._whitelist).calculate_operations()
+                solution = solver.solve(use_latest=self._whitelist)
+                solved_packages = solution.get_solved_packages()
+
+            if not self.executor.enabled:
+                # If we are only in lock mode, no need to go any further
+                self._write_lock_file(solved_packages)
+                return 0
 
             lockfile_repo = LockfileRepository()
-            self._populate_lockfile_repo(lockfile_repo, ops)
+            for package in solved_packages:
+                if not lockfile_repo.has_package(package):
+                    lockfile_repo.add_package(package)
 
         else:
             self._io.write_line("<info>Installing dependencies from lock file</>")
@@ -261,11 +266,6 @@ def _do_install(self) -> int:
             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._io.is_verbose():
             self._io.write_line("")
             self._io.write_line(
@@ -312,13 +312,17 @@ def _do_install(self) -> int:
 
         if status == 0 and self._update:
             # Only write lock file when installation is success
-            self._write_lock_file(lockfile_repo)
+            self._write_lock_file(solved_packages)
 
         return status
 
-    def _write_lock_file(self, repo: LockfileRepository, force: bool = False) -> None:
+    def _write_lock_file(
+        self,
+        packages: dict[Package, TransitivePackageInfo],
+        force: bool = False,
+    ) -> None:
         if not self.is_dry_run() and (force or self._update):
-            updated_lock = self._locker.set_lock_data(self._package, repo.packages)
+            updated_lock = self._locker.set_lock_data(self._package, packages)
 
             if updated_lock:
                 self._io.write_line("")
@@ -327,16 +331,5 @@ def _write_lock_file(self, repo: LockfileRepository, force: bool = False) -> Non
     def _execute(self, operations: list[Operation]) -> int:
         return self._executor.execute(operations)
 
-    def _populate_lockfile_repo(
-        self, repo: LockfileRepository, ops: Iterable[Operation]
-    ) -> None:
-        for op in ops:
-            if isinstance(op, Uninstall):
-                continue
-
-            package = op.target_package if isinstance(op, Update) else op.package
-            if not repo.has_package(package):
-                repo.add_package(package)
-
     def _get_installed(self) -> InstalledRepository:
         return InstalledRepository.load(self._env)

src/poetry/packages/locker.py

@@ -38,6 +38,7 @@
     from poetry.core.packages.vcs_dependency import VCSDependency
     from tomlkit.toml_document import TOMLDocument
 
+    from poetry.packages.transitive_package_info import TransitivePackageInfo
     from poetry.repositories.lockfile_repository import LockfileRepository
 
 logger = logging.getLogger(__name__)
@@ -49,7 +50,7 @@
 
 
 class Locker:
-    _VERSION = "2.0"
+    _VERSION = "2.1"
     _READ_VERSION_RANGE = ">=1,<3"
 
     _legacy_keys: ClassVar[list[str]] = [
@@ -231,7 +232,9 @@ def locked_repository(self) -> LockfileRepository:
 
         return repository
 
-    def set_lock_data(self, root: Package, packages: list[Package]) -> bool:
+    def set_lock_data(
+        self, root: Package, packages: dict[Package, TransitivePackageInfo]
+    ) -> bool:
         """Store lock data and eventually persist to the lock file"""
         lock = self._compute_lock_data(root, packages)
 
@@ -242,7 +245,7 @@ def set_lock_data(self, root: Package, packages: list[Package]) -> bool:
         return False
 
     def _compute_lock_data(
-        self, root: Package, packages: list[Package]
+        self, root: Package, packages: dict[Package, TransitivePackageInfo]
     ) -> TOMLDocument:
         package_specs = self._lock_packages(packages)
         # Retrieving hashes
@@ -356,7 +359,9 @@ def _get_lock_data(self) -> dict[str, Any]:
 
         return lock_data
 
-    def _lock_packages(self, packages: list[Package]) -> list[dict[str, Any]]:
+    def _lock_packages(
+        self, packages: dict[Package, TransitivePackageInfo]
+    ) -> list[dict[str, Any]]:
         locked = []
 
         for package in sorted(
@@ -371,13 +376,15 @@ def _lock_packages(self, packages: list[Package]) -> list[dict[str, Any]]:
                 x.source_resolved_reference or "",
             ),
         ):
-            spec = self._dump_package(package)
+            spec = self._dump_package(package, packages[package])
 
             locked.append(spec)
 
         return locked
 
-    def _dump_package(self, package: Package) -> dict[str, Any]:
+    def _dump_package(
+        self, package: Package, transitive_info: TransitivePackageInfo
+    ) -> dict[str, Any]:
         dependencies: dict[str, list[Any]] = {}
         for dependency in sorted(
             package.requires,
@@ -447,8 +454,21 @@ def _dump_package(self, package: Package) -> dict[str, Any]:
             "description": package.description or "",
             "optional": package.optional,
             "python-versions": package.python_versions,
-            "files": sorted(package.files, key=lambda x: x["file"]),
+            "groups": sorted(transitive_info.groups, key=lambda x: (x != "main", x)),
         }
+        if transitive_info.markers:
+            if len(markers := set(transitive_info.markers.values())) == 1:
+                if not (marker := next(iter(markers))).is_any():
+                    data["markers"] = str(marker)
+            else:
+                data["markers"] = inline_table()
+                for k, v in sorted(
+                    transitive_info.markers.items(),
+                    key=lambda x: (x[0] != "main", x[0]),
+                ):
+                    if not v.is_any():
+                        data["markers"][k] = str(v)
+        data["files"] = sorted(package.files, key=lambda x: x["file"])
 
         if dependencies:
             data["dependencies"] = table()

src/poetry/packages/transitive_package_info.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from poetry.core.version.markers import BaseMarker
+
+
+@dataclass
+class TransitivePackageInfo:
+    depth: int  # max depth in the dependency tree
+    groups: set[str]
+    markers: dict[str, BaseMarker]  # group -> marker

src/poetry/puzzle/provider.py

@@ -525,6 +525,7 @@ def complete_package(
             package = dependency_package.package
             dependency = dependency_package.dependency
             new_dependency = package.without_features().to_dependency()
+            new_dependency.marker = AnyMarker()
 
             # When adding dependency foo[extra] -> foo, preserve foo's source, if it's
             # specified. This prevents us from trying to get foo from PyPI