Add pylock select function

src/packaging/pylock.py

@@ -12,11 +12,13 @@
     Callable,
     Protocol,
     TypeVar,
+    cast,
 )
 from urllib.parse import urlparse
 
-from .markers import Marker
+from .markers import Environment, Marker, default_environment
 from .specifiers import SpecifierSet
+from .tags import sys_tags
 from .utils import (
     NormalizedName,
     is_normalized_name,
@@ -26,10 +28,13 @@
 from .version import Version
 
 if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Collection, Iterator
     from pathlib import Path
 
     from typing_extensions import Self
 
+    from .tags import Tag
+
 _logger = logging.getLogger(__name__)
 
 __all__ = [
@@ -300,6 +305,10 @@ class PylockUnsupportedVersionError(PylockValidationError):
     """Raised when encountering an unsupported `lock_version`."""
 
 
+class PylockSelectError(Exception):
+    """Base exception for errors raised by :meth:`Pylock.select`."""
+
+
 @dataclass(frozen=True, init=False)
 class PackageVcs:
     type: str
@@ -692,3 +701,173 @@ def validate(self) -> None:
 
         Raises :class:`PylockValidationError` otherwise."""
         self.from_dict(self.to_dict())
+
+    def select(
+        self,
+        *,
+        environment: Environment | None = None,
+        tags: Sequence[Tag] | None = None,
+        extras: Collection[str] | None = None,
+        dependency_groups: Collection[str] | None = None,
+    ) -> Iterator[
+        tuple[
+            Package,
+            PackageVcs
+            | PackageDirectory
+            | PackageArchive
+            | PackageWheel
+            | PackageSdist,
+        ]
+    ]:
+        """Select what to install from the lock file.
+
+        The *environment* and *tags* parameters represent the environment being
+        selected for. If unspecified,
+        ``packaging.markers.default_environment()`` and
+        ``packaging.tags.sys_tags()`` are used.
+
+        The *extras* parameter represents the extras to install.
+
+        The *dependency_groups* parameter represents the groups to install. If
+        unspecified, the default groups are used.
+
+        For better error reporting, it is recommended to use this method on
+        valid Pylock instances (i.e. one obtained from :meth:`Pylock.from_dict`
+        or if constructed manually, after calling :meth:`Pylock.validate`).
+        """
+        supported_tags = frozenset(tags or sys_tags())
+
+        # #. Gather the extras and dependency groups to install and set ``extras`` and
+        #    ``dependency_groups`` for marker evaluation, respectively.
+        #
+        #    #. ``extras`` SHOULD be set to the empty set by default.
+        #    #. ``dependency_groups`` SHOULD be the set created from
+        #       :ref:`pylock-default-groups` by default.
+        env = cast(
+            "dict[str, str | frozenset[str]]",
+            dict(
+                environment or {},  # Marker.evaluate will fill-up
+                extras=frozenset(extras or []),
+                dependency_groups=frozenset(
+                    (self.default_groups or [])
+                    if dependency_groups is None  # to allow selecting no group
+                    else dependency_groups
+                ),
+            ),
+        )
+        env_python_version = (
+            environment["python_version"]
+            if environment
+            else default_environment()["python_version"]
+        )
+
+        # #. Check if the metadata version specified by :ref:`pylock-lock-version` is
+        #    supported; an error or warning MUST be raised as appropriate.
+        # Covered by lock.validate() which is a precondition for this method.
+
+        # #. If :ref:`pylock-requires-python` is specified, check that the environment
+        #    being installed for meets the requirement; an error MUST be raised if it is
+        #    not met.
+        if self.requires_python and not self.requires_python.contains(
+            env_python_version,
+        ):
+            raise PylockSelectError(
+                f"Provided environment does not satisfy the Python version "
+                f"requirement {self.requires_python!r}"
+            )
+
+        # #. If :ref:`pylock-environments` is specified, check that at least one of the
+        #    environment marker expressions is satisfied; an error MUST be raised if no
+        #    expression is satisfied.
+        if self.environments:
+            for env_marker in self.environments:
+                if env_marker.evaluate(
+                    cast("dict[str, str]", environment or {}), context="requirement"
+                ):
+                    break
+            else:
+                raise PylockSelectError(
+                    "Provided environment does not satisfy any of the "
+                    "environments specified in the lock file"
+                )
+
+        # #. For each package listed in :ref:`pylock-packages`:
+        selected_packages_by_name: dict[str, tuple[int, Package]] = {}
+        for package_index, package in enumerate(self.packages):
+            # #. If :ref:`pylock-packages-marker` is specified, check if it is
+            #    satisfied;if it isn't, skip to the next package.
+            if package.marker and not package.marker.evaluate(env, context="lock_file"):
+                continue
+
+            # #. If :ref:`pylock-packages-requires-python` is specified, check if it is
+            #    satisfied; an error MUST be raised if it isn't.
+            if package.requires_python and not package.requires_python.contains(
+                env_python_version,
+            ):
+                raise PylockSelectError(
+                    f"Provided environment does not satisfy the Python version "
+                    f"requirement {package.requires_python!r} for package "
+                    f"{package.name!r} at packages[{package_index}]"
+                )
+
+            # #. Check that no other conflicting instance of the package has been slated
+            #    to be installed; an error about the ambiguity MUST be raised otherwise.
+            if package.name in selected_packages_by_name:
+                raise PylockSelectError(
+                    f"Multiple packages with the name {package.name!r} are "
+                    f"selected at packages[{package_index}] and "
+                    f"packages[{selected_packages_by_name[package.name][0]}]"
+                )
+
+            # #. Check that the source of the package is specified appropriately (i.e.
+            #    there are no conflicting sources in the package entry);
+            #    an error MUST be raised if any issues are found.
+            # Covered by lock.validate() which is a precondition for this method.
+
+            # #. Add the package to the set of packages to install.
+            selected_packages_by_name[package.name] = (package_index, package)
+
+        # #. For each package to be installed:
+        for package_index, package in selected_packages_by_name.values():
+            # - If :ref:`pylock-packages-vcs` is set:
+            if package.vcs is not None:
+                yield package, package.vcs
+
+            # - Else if :ref:`pylock-packages-directory` is set:
+            elif package.directory is not None:
+                yield package, package.directory
+
+            # - Else if :ref:`pylock-packages-archive` is set:
+            elif package.archive is not None:
+                yield package, package.archive
+
+            # - Else if there are entries for :ref:`pylock-packages-wheels`:
+            elif package.wheels:
+                # #. Look for the appropriate wheel file based on
+                #    :ref:`pylock-packages-wheels-name`; if one is not found then move
+                #    on to :ref:`pylock-packages-sdist` or an error MUST be raised about
+                #    a lack of source for the project.
+                for wheel in package.wheels:
+                    wheel_tags = parse_wheel_filename(wheel.filename)[-1]
+                    if not wheel_tags.isdisjoint(supported_tags):
+                        yield package, wheel
+                        break
+                else:
+                    if package.sdist is not None:
+                        yield package, package.sdist
+                    else:
+                        raise PylockSelectError(
+                            f"No wheel found matching the provided tags "
+                            f"for package {package.name!r} "
+                            f"at packages[{package_index}], "
+                            f"and no sdist available as a fallback"
+                        )
+
+            # - Else if no :ref:`pylock-packages-wheels` file is found or
+            #   :ref:`pylock-packages-sdist` is solely set:
+            elif package.sdist is not None:
+                yield package, package.sdist
+
+            else:
+                # Covered by lock.validate() which is a precondition for this method.
+                raise NotImplementedError  # pragma: no cover