feat: support 'extras' and 'dependency_groups' markers (#11)

Signed-off-by: Frost Ming <me@frostming.com>

Author: frostming

src/dep_logic/markers/__init__.py

@@ -31,15 +31,15 @@
 
 
 __all__ = [
-    "parse_marker",
-    "from_pkg_marker",
-    "InvalidMarker",
-    "BaseMarker",
     "AnyMarker",
+    "BaseMarker",
     "EmptyMarker",
+    "InvalidMarker",
     "MarkerExpression",
     "MarkerUnion",
     "MultiMarker",
+    "from_pkg_marker",
+    "parse_marker",
 ]
 
 
@@ -99,3 +99,34 @@ def _build_markers(markers: _ParsedMarkers) -> BaseMarker:
         else:
             or_groups[-1] &= _build_markers(item)
     return MarkerUnion.of(*or_groups)
+
+
+def _patch_marker_parser() -> None:
+    import re
+
+    try:
+        from packaging._tokenizer import DEFAULT_RULES
+    except (ModuleNotFoundError, AttributeError):
+        return
+
+    DEFAULT_RULES["VARIABLE"] = re.compile(
+        r"""
+            \b(
+                python_version
+                |python_full_version
+                |os[._]name
+                |sys[._]platform
+                |platform_(release|system)
+                |platform[._](version|machine|python_implementation)
+                |python_implementation
+                |implementation_(name|version)
+                |extras?
+                |dependency_groups
+            )\b
+        """,
+        re.VERBOSE,
+    )
+
+
+_patch_marker_parser()
+del _patch_marker_parser

src/dep_logic/markers/any.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dep_logic.markers.base import BaseMarker
+from dep_logic.markers.base import BaseMarker, EvaluationContext
 
 
 class AnyMarker(BaseMarker):
@@ -17,7 +17,11 @@ def __or__(self, other: BaseMarker) -> BaseMarker:
     def is_any(self) -> bool:
         return True
 
-    def evaluate(self, environment: dict[str, str] | None = None) -> bool:
+    def evaluate(
+        self,
+        environment: dict[str, str | set[str]] | None = None,
+        context: EvaluationContext = "metadata",
+    ):
         return True
 
     def without_extras(self) -> BaseMarker:

src/dep_logic/markers/base.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 from abc import ABCMeta, abstractmethod
-from typing import Any
+from typing import Any, Literal
+
+EvaluationContext = Literal["lock_file", "metadata", "requirement"]
 
 
 class BaseMarker(metaclass=ABCMeta):
@@ -28,7 +30,11 @@ def is_empty(self) -> bool:
         return False
 
     @abstractmethod
-    def evaluate(self, environment: dict[str, str] | None = None) -> bool:
+    def evaluate(
+        self,
+        environment: dict[str, str | set[str]] | None = None,
+        context: EvaluationContext = "metadata",
+    ) -> bool:
         raise NotImplementedError
 
     @abstractmethod

src/dep_logic/markers/empty.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dep_logic.markers.base import BaseMarker
+from dep_logic.markers.base import BaseMarker, EvaluationContext
 
 
 class EmptyMarker(BaseMarker):
@@ -17,7 +17,11 @@ def __or__(self, other: BaseMarker) -> BaseMarker:
     def is_empty(self) -> bool:
         return True
 
-    def evaluate(self, environment: dict[str, str] | None = None) -> bool:
+    def evaluate(
+        self,
+        environment: dict[str, str | set[str]] | None = None,
+        context: EvaluationContext = "metadata",
+    ) -> bool:
         return False
 
     def without_extras(self) -> BaseMarker:

src/dep_logic/markers/multi.py

@@ -4,7 +4,7 @@
 from typing import Iterator
 
 from dep_logic.markers.any import AnyMarker
-from dep_logic.markers.base import BaseMarker
+from dep_logic.markers.base import BaseMarker, EvaluationContext
 from dep_logic.markers.empty import EmptyMarker
 from dep_logic.markers.single import MarkerExpression, SingleMarker
 from dep_logic.utils import DATACLASS_ARGS, flatten_items, intersection, union
@@ -135,8 +135,12 @@ def union_simplify(self, other: BaseMarker) -> BaseMarker | None:
 
         return None
 
-    def evaluate(self, environment: dict[str, str] | None = None) -> bool:
-        return all(m.evaluate(environment) for m in self.markers)
+    def evaluate(
+        self,
+        environment: dict[str, str | set[str]] | None = None,
+        context: EvaluationContext = "metadata",
+    ) -> bool:
+        return all(m.evaluate(environment, context) for m in self.markers)
 
     def without_extras(self) -> BaseMarker:
         return self.exclude("extra")

src/dep_logic/markers/single.py

@@ -1,24 +1,44 @@
 from __future__ import annotations
 
 import functools
+import operator
 import typing as t
+from abc import abstractmethod
 from dataclasses import dataclass, field, replace
 
-from packaging.markers import Marker as _Marker
+from packaging.markers import default_environment
+from packaging.specifiers import InvalidSpecifier, Specifier
+from packaging.version import InvalidVersion
 
 from dep_logic.markers.any import AnyMarker
-from dep_logic.markers.base import BaseMarker
+from dep_logic.markers.base import BaseMarker, EvaluationContext
 from dep_logic.markers.empty import EmptyMarker
 from dep_logic.specifiers import BaseSpecifier
 from dep_logic.specifiers.base import VersionSpecifier
 from dep_logic.specifiers.generic import GenericSpecifier
-from dep_logic.utils import DATACLASS_ARGS, OrderedSet, get_reflect_op
+from dep_logic.utils import DATACLASS_ARGS, OrderedSet, get_reflect_op, normalize_name
 
 if t.TYPE_CHECKING:
     from dep_logic.markers.multi import MultiMarker
     from dep_logic.markers.union import MarkerUnion
 
 PYTHON_VERSION_MARKERS = {"python_version", "python_full_version"}
+MARKERS_ALLOWING_SET = {"extras", "dependency_groups"}
+Operator = t.Callable[[str, t.Union[str, t.Set[str]]], bool]
+_operators: dict[str, Operator] = {
+    "in": lambda lhs, rhs: lhs in rhs,
+    "not in": lambda lhs, rhs: lhs not in rhs,
+    "<": operator.lt,
+    "<=": operator.le,
+    "==": operator.eq,
+    "!=": operator.ne,
+    ">=": operator.ge,
+    ">": operator.gt,
+}
+
+
+class UndefinedComparison(ValueError):
+    pass
 
 
 class SingleMarker(BaseMarker):
@@ -44,16 +64,25 @@ def only(self, *marker_names: str) -> BaseMarker:
 
         return self
 
-    def evaluate(self, environment: dict[str, str] | None = None) -> bool:
-        pkg_marker = _Marker(str(self))
-        if self.name != "extra" or not environment or "extra" not in environment:
-            return pkg_marker.evaluate(environment)
-        extras = [extra] if isinstance(extra := environment["extra"], str) else extra
-        assert isinstance(self, MarkerExpression)
-        is_negated = self.op in ("not in", "!=")
-        if is_negated:
-            return all(pkg_marker.evaluate({"extra": extra}) for extra in extras)
-        return any(pkg_marker.evaluate({"extra": extra}) for extra in extras)
+    def evaluate(
+        self,
+        environment: dict[str, str | set[str]] | None = None,
+        context: EvaluationContext = "metadata",
+    ) -> bool:
+        current_environment = t.cast("dict[str, str|set[str]]", default_environment())
+        if context == "metadata":
+            current_environment["extra"] = ""
+        elif context == "lock_file":
+            current_environment.update(extras=set(), dependency_groups=set())
+        if environment:
+            current_environment.update(environment)
+        if "extra" in current_environment and current_environment["extra"] is None:
+            current_environment["extra"] = ""
+        return self._evaluate(current_environment)
+
+    @abstractmethod
+    def _evaluate(self, environment: dict[str, str | set[str]]) -> bool:
+        raise NotImplementedError
 
 
 @dataclass(unsafe_hash=True, **DATACLASS_ARGS)
@@ -141,6 +170,46 @@ def __or__(self, other: t.Any) -> BaseMarker:
 
         return MarkerUnion(self, other)
 
+    def _evaluate(self, environment: dict[str, str | set[str]]) -> bool:
+        if self.name == "extra":
+            # Support batch comparison for "extra" markers
+            extra = environment["extra"]
+            if isinstance(extra, str):
+                extra = {extra}
+            assert self.op in ("==", "!=")
+            value = normalize_name(self.value)
+            extra = {normalize_name(v) for v in extra}
+            return value in extra if self.op == "==" else value not in extra
+
+        target = environment[self.name]
+        if self.reversed:
+            lhs, rhs = self.value, target
+            oper = _operators.get(get_reflect_op(self.op))
+        else:
+            lhs, rhs = target, self.value
+            assert isinstance(lhs, str)
+            oper = _operators.get(self.op)
+        if self.name in MARKERS_ALLOWING_SET:
+            lhs = normalize_name(lhs)
+            if isinstance(rhs, set):
+                rhs = {normalize_name(v) for v in rhs}
+            else:
+                rhs = normalize_name(rhs)
+        if isinstance(rhs, str):
+            try:
+                spec = Specifier(f"{self.op}{rhs}")
+            except InvalidSpecifier:
+                pass
+            else:
+                try:
+                    return spec.contains(lhs)
+                except InvalidVersion:
+                    pass
+
+        if oper is None:
+            raise UndefinedComparison(f"Undefined comparison {self}")
+        return oper(lhs, rhs)
+
 
 @dataclass(frozen=True, unsafe_hash=True, **DATACLASS_ARGS)
 class EqualityMarkerUnion(SingleMarker):
@@ -210,6 +279,9 @@ def __or__(self, other: t.Any) -> BaseMarker:
     __rand__ = __and__
     __ror__ = __or__
 
+    def _evaluate(self, environment: dict[str, str | set[str]]) -> bool:
+        return environment[self.name] in self.values
+
 
 @dataclass(frozen=True, unsafe_hash=True, **DATACLASS_ARGS)
 class InequalityMultiMarker(SingleMarker):
@@ -283,6 +355,9 @@ def __or__(self, other: t.Any) -> BaseMarker:
     __rand__ = __and__
     __ror__ = __or__
 
+    def _evaluate(self, environment: dict[str, str | set[str]]) -> bool:
+        return environment[self.name] not in self.values
+
 
 @functools.lru_cache(maxsize=None)
 def _merge_single_markers(
@@ -375,5 +450,5 @@ def _normalize_python_version_specifier(marker: MarkerExpression) -> BaseSpecifi
         splitted[-1] = str(int(splitted[-1]) + 1)
         op = "<"
 
-    spec = parse_version_specifier(f'{op}{".".join(splitted)}')
+    spec = parse_version_specifier(f"{op}{'.'.join(splitted)}")
     return spec

src/dep_logic/markers/union.py

@@ -4,7 +4,7 @@
 from typing import Iterator
 
 from dep_logic.markers.any import AnyMarker
-from dep_logic.markers.base import BaseMarker
+from dep_logic.markers.base import BaseMarker, EvaluationContext
 from dep_logic.markers.empty import EmptyMarker
 from dep_logic.markers.multi import MultiMarker
 from dep_logic.markers.single import SingleMarker
@@ -135,8 +135,12 @@ def intersect_simplify(self, other: BaseMarker) -> BaseMarker | None:
 
         return None
 
-    def evaluate(self, environment: dict[str, str] | None = None) -> bool:
-        return any(m.evaluate(environment) for m in self.markers)
+    def evaluate(
+        self,
+        environment: dict[str, str | set[str]] | None = None,
+        context: EvaluationContext = "metadata",
+    ) -> bool:
+        return any(m.evaluate(environment, context) for m in self.markers)
 
     def without_extras(self) -> BaseMarker:
         return self.exclude("extra")

src/dep_logic/utils.py

@@ -191,3 +191,7 @@ def __len__(self) -> int:
 
     def peek(self) -> T:
         return self._data[0]
+
+
+def normalize_name(name: str) -> str:
+    return re.sub(r"[-_.]+", "-", name).lower()