Add pylock select function

src/packaging/pylock.py

@@ -6,24 +6,23 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Protocol,
-    TypeVar,
-)
-
-from .markers import Marker
+from typing import TYPE_CHECKING, Any, Callable, Protocol, TypeVar, cast
+
+from .markers import Marker, default_environment
 from .specifiers import SpecifierSet
-from .utils import NormalizedName, is_normalized_name
+from .tags import sys_tags
+from .utils import NormalizedName, is_normalized_name, parse_wheel_filename
 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 .markers import Environment
+    from .tags import Tag
+
 _logger = logging.getLogger(__name__)
 
 __all__ = [
@@ -274,6 +273,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
@@ -638,3 +641,181 @@ 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[  # XXX or Iterable?
+        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`).
+        """
+        if environment is None:
+            environment = default_environment()
+        if tags is None:
+            tags = list(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: dict[str, str | frozenset[str]] = {
+            **cast("dict[str, str]", environment),
+            "extras": frozenset(extras or []),
+            "dependency_groups": frozenset(
+                dependency_groups or self.default_groups or []
+            ),
+        }
+        env_python_version = environment.get("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() above.
+
+        # #. 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 is not None:
+            if not env_python_version:
+                raise PylockSelectError(
+                    f"Provided environment does not specify a Python version, "
+                    f"but the lock file requires Python {self.requires_python!r}"
+                )
+            if not self.requires_python.contains(env_python_version, prereleases=True):
+                # XXX confirm prereleases=True
+                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(env, context="lock_file"):  # XXX check context
+                    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="requirement"
+            ):  # XXX check context
+                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:
+                if not env_python_version:
+                    raise PylockSelectError(
+                        f"Provided environment does not specify a Python version, "
+                        f"but package {package.name!r} at packages[{package_index}] "
+                        f"requires Python {package.requires_python!r}"
+                    )
+                if not package.requires_python.contains(
+                    env_python_version, prereleases=True
+                ):
+                    # XXX confirm prereleases=True
+                    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() above.
+
+            # #. 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 package_wheel in package.wheels:
+                    assert package_wheel.name  # XXX get name from path or url
+                    package_wheel_tags = parse_wheel_filename(package_wheel.name)[-1]
+                    if not package_wheel_tags.isdisjoint(tags):
+                        yield package, 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() above.
+                raise NotImplementedError

tests/test_pylock.py

@@ -2,24 +2,27 @@
 
 import datetime
 import sys
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 import pytest
 import tomli_w
 
-from packaging.markers import Marker
+from packaging.markers import Environment, Marker
 from packaging.pylock import (
     Package,
     PackageDirectory,
     PackageVcs,
     PackageWheel,
     Pylock,
+    PylockSelectError,
     PylockUnsupportedVersionError,
     PylockValidationError,
     is_valid_pylock_path,
 )
 from packaging.specifiers import SpecifierSet
+from packaging.tags import Tag
 from packaging.utils import NormalizedName
 from packaging.version import Version
 
@@ -578,3 +581,60 @@ def test_validate_attestation_identity_invalid_kind() -> None:
         "Unexpected type int (expected str) "
         "in 'packages[0].attestation-identities[0].kind'"
     )
+
+
+@dataclass
+class Platform:
+    tags: list[Tag]
+    environment: Environment
+
+
+_py312_linux = Platform(
+    tags=[
+        Tag("cp312", "cp312", "manylinux_2_17_x86_64"),
+        Tag("py3", "none", "any"),
+    ],
+    environment={
+        "implementation_name": "cpython",
+        "implementation_version": "3.12.12",
+        "os_name": "posix",
+        "platform_machine": "x86_64",
+        "platform_release": "6.8.0-100-generic",
+        "platform_system": "Linux",
+        "platform_version": "#100-Ubuntu SMP PREEMPT_DYNAMIC",
+        "python_full_version": "3.12.12",
+        "platform_python_implementation": "CPython",
+        "python_version": "3.12",
+        "sys_platform": "linux",
+    },
+)
+
+
+def test_select_smoke_test() -> None:
+    pylock_path = Path(__file__).parent / "pylock" / "pylock.spec-example.toml"
+    lock = Pylock.from_dict(tomllib.loads(pylock_path.read_text()))
+    for package, dist in lock.select(
+        tags=_py312_linux.tags,
+        environment=_py312_linux.environment,
+    ):
+        assert isinstance(package, Package)
+        assert isinstance(dist, PackageWheel)
+
+
+def test_select_require_python_mismatch() -> None:
+    pylock = Pylock(
+        lock_version=Version("1.0"),
+        created_by="some_tool",
+        requires_python=SpecifierSet("==3.14.*"),
+        packages=[],
+    )
+    with pytest.raises(
+        PylockSelectError,
+        match="Provided environment does not satisfy the Python version requirement",
+    ):
+        list(
+            pylock.select(
+                tags=_py312_linux.tags,
+                environment=_py312_linux.environment,
+            )
+        )