support 'extras' and 'dependency_groups' markers (#888)

Signed-off-by: Frost Ming <me@frostming.com>
Co-authored-by: Brett Cannon <brett@python.org>

Author: frostming

CHANGELOG.rst

@@ -4,7 +4,7 @@ Changelog
 *unreleased*
 ~~~~~~~~~~~~
 
-No unreleased changes.
+* PEP 751: Add support for ``extras`` and ``dependency_groups`` markers. (:issue:`885`)
 
 24.2 - 2024-11-08
 ~~~~~~~~~~~~~~~~~

docs/markers.rst

@@ -64,6 +64,9 @@ Reference
 
     :param dict environment: A dictionary containing keys and values to
                              override the detected environment.
+    :param str context: A string representing the context in which the marker is evaluated.
+                        Acceptable values are "metadata" (for core metadata; default),
+                        "lock_file", and "requirement" (i.e. all other situations).
     :raises: UndefinedComparison: If the marker uses a comparison on strings
                                   which are not valid versions per the
                                   :ref:`specification of version specifiers

src/packaging/_tokenizer.py

@@ -68,7 +68,8 @@ def __str__(self) -> str:
                 |platform[._](version|machine|python_implementation)
                 |python_implementation
                 |implementation_(name|version)
-                |extra
+                |extras?
+                |dependency_groups
             )\b
         """,
         re.VERBOSE,

src/packaging/markers.py

@@ -8,7 +8,7 @@
 import os
 import platform
 import sys
-from typing import Any, Callable, TypedDict, cast
+from typing import AbstractSet, Any, Callable, Literal, TypedDict, Union, cast
 
 from ._parser import MarkerAtom, MarkerList, Op, Value, Variable
 from ._parser import parse_marker as _parse_marker
@@ -17,14 +17,17 @@
 from .utils import canonicalize_name
 
 __all__ = [
+    "EvaluateContext",
     "InvalidMarker",
     "Marker",
     "UndefinedComparison",
     "UndefinedEnvironmentName",
     "default_environment",
 ]
 
-Operator = Callable[[str, str], bool]
+Operator = Callable[[str, Union[str, AbstractSet[str]]], bool]
+EvaluateContext = Literal["metadata", "lock_file", "requirement"]
+MARKERS_ALLOWING_SET = {"extras", "dependency_groups"}
 
 
 class InvalidMarker(ValueError):
@@ -174,13 +177,14 @@ def _format_marker(
 }
 
 
-def _eval_op(lhs: str, op: Op, rhs: str) -> bool:
-    try:
-        spec = Specifier("".join([op.serialize(), rhs]))
-    except InvalidSpecifier:
-        pass
-    else:
-        return spec.contains(lhs, prereleases=True)
+def _eval_op(lhs: str, op: Op, rhs: str | AbstractSet[str]) -> bool:
+    if isinstance(rhs, str):
+        try:
+            spec = Specifier("".join([op.serialize(), rhs]))
+        except InvalidSpecifier:
+            pass
+        else:
+            return spec.contains(lhs, prereleases=True)
 
     oper: Operator | None = _operators.get(op.serialize())
     if oper is None:
@@ -189,19 +193,29 @@ def _eval_op(lhs: str, op: Op, rhs: str) -> bool:
     return oper(lhs, rhs)
 
 
-def _normalize(*values: str, key: str) -> tuple[str, ...]:
+def _normalize(
+    lhs: str, rhs: str | AbstractSet[str], key: str
+) -> tuple[str, str | AbstractSet[str]]:
     # PEP 685 – Comparison of extra names for optional distribution dependencies
     # https://peps.python.org/pep-0685/
     # > When comparing extra names, tools MUST normalize the names being
     # > compared using the semantics outlined in PEP 503 for names
     if key == "extra":
-        return tuple(canonicalize_name(v) for v in values)
+        assert isinstance(rhs, str), "extra value must be a string"
+        return (canonicalize_name(lhs), canonicalize_name(rhs))
+    if key in MARKERS_ALLOWING_SET:
+        if isinstance(rhs, str):  # pragma: no cover
+            return (canonicalize_name(lhs), canonicalize_name(rhs))
+        else:
+            return (canonicalize_name(lhs), {canonicalize_name(v) for v in rhs})
 
     # other environment markers don't have such standards
-    return values
+    return lhs, rhs
 
 
-def _evaluate_markers(markers: MarkerList, environment: dict[str, str]) -> bool:
+def _evaluate_markers(
+    markers: MarkerList, environment: dict[str, str | AbstractSet[str]]
+) -> bool:
     groups: list[list[bool]] = [[]]
 
     for marker in markers:
@@ -220,7 +234,7 @@ def _evaluate_markers(markers: MarkerList, environment: dict[str, str]) -> bool:
                 lhs_value = lhs.value
                 environment_key = rhs.value
                 rhs_value = environment[environment_key]
-
+            assert isinstance(lhs_value, str), "lhs must be a string"
             lhs_value, rhs_value = _normalize(lhs_value, rhs_value, key=environment_key)
             groups[-1].append(_eval_op(lhs_value, op, rhs_value))
         else:
@@ -298,34 +312,51 @@ def __eq__(self, other: Any) -> bool:
 
         return str(self) == str(other)
 
-    def evaluate(self, environment: dict[str, str] | None = None) -> bool:
+    def evaluate(
+        self,
+        environment: dict[str, str] | None = None,
+        context: EvaluateContext = "metadata",
+    ) -> bool:
         """Evaluate a marker.
 
         Return the boolean from evaluating the given marker against the
         environment. environment is an optional argument to override all or
-        part of the determined environment.
+        part of the determined environment. The *context* parameter specifies what
+        context the markers are being evaluated for, which influences what markers
+        are considered valid. Acceptable values are "metadata" (for core metadata;
+        default), "lock_file", and "requirement" (i.e. all other situations).
 
         The environment is determined from the current Python process.
         """
-        current_environment = cast("dict[str, str]", default_environment())
-        current_environment["extra"] = ""
+        current_environment = cast(
+            "dict[str, str | AbstractSet[str]]", default_environment()
+        )
+        if context == "lock_file":
+            current_environment.update(
+                extras=frozenset(), dependency_groups=frozenset()
+            )
+        elif context == "metadata":
+            current_environment["extra"] = ""
         if environment is not None:
             current_environment.update(environment)
             # The API used to allow setting extra to None. We need to handle this
             # case for backwards compatibility.
-            if current_environment["extra"] is None:
+            if "extra" in current_environment and current_environment["extra"] is None:
                 current_environment["extra"] = ""
 
         return _evaluate_markers(
             self._markers, _repair_python_full_version(current_environment)
         )
 
 
-def _repair_python_full_version(env: dict[str, str]) -> dict[str, str]:
+def _repair_python_full_version(
+    env: dict[str, str | AbstractSet[str]],
+) -> dict[str, str | AbstractSet[str]]:
     """
     Work around platform.python_version() returning something that is not PEP 440
     compliant for non-tagged Python builds.
     """
-    if env["python_full_version"].endswith("+"):
-        env["python_full_version"] += "local"
+    python_full_version = cast(str, env["python_full_version"])
+    if python_full_version.endswith("+"):
+        env["python_full_version"] = f"{python_full_version}local"
     return env

tests/test_markers.py

@@ -394,3 +394,35 @@ def test_python_full_version_untagged_user_provided(self):
     def test_python_full_version_untagged(self):
         with mock.patch("platform.python_version", return_value="3.11.1+"):
             assert Marker("python_full_version < '3.12'").evaluate()
+
+    @pytest.mark.parametrize("variable", ["extras", "dependency_groups"])
+    @pytest.mark.parametrize(
+        "expression,result",
+        [
+            pytest.param('"foo" in {0}', True, id="value-in-foo"),
+            pytest.param('"bar" in {0}', True, id="value-in-bar"),
+            pytest.param('"baz" in {0}', False, id="value-not-in"),
+            pytest.param('"baz" not in {0}', True, id="value-not-in-negated"),
+            pytest.param('"foo" in {0} and "bar" in {0}', True, id="and-in"),
+            pytest.param('"foo" in {0} or "bar" in {0}', True, id="or-in"),
+            pytest.param(
+                '"baz" in {0} and "foo" in {0}', False, id="short-circuit-and"
+            ),
+            pytest.param('"foo" in {0} or "baz" in {0}', True, id="short-circuit-or"),
+            pytest.param('"Foo" in {0}', True, id="case-sensitive"),
+        ],
+    )
+    def test_extras_and_dependency_groups(self, variable, expression, result):
+        environment = {variable: {"foo", "bar"}}
+        assert Marker(expression.format(variable)).evaluate(environment) == result
+
+    @pytest.mark.parametrize("variable", ["extras", "dependency_groups"])
+    def test_extras_and_dependency_groups_disallowed(self, variable):
+        marker = Marker(f'"foo" in {variable}')
+        assert not marker.evaluate(context="lock_file")
+
+        with pytest.raises(KeyError):
+            marker.evaluate()
+
+        with pytest.raises(KeyError):
+            marker.evaluate(context="requirement")