constraints: implement intersection and union for extra constraints

Extra constraints are slightly different from standard generic constraints. For example, the standard generic constraint "==linux, ==win32" can never be satisfied (i.e., is "empty"), but the extra constraint "==extra1, ==extra2" can be satisfied (if both extras are requested).

src/poetry/core/constraints/generic/__init__.py

@@ -6,6 +6,7 @@
 from poetry.core.constraints.generic.empty_constraint import EmptyConstraint
 from poetry.core.constraints.generic.multi_constraint import MultiConstraint
 from poetry.core.constraints.generic.parser import parse_constraint
+from poetry.core.constraints.generic.parser import parse_extra_constraint
 from poetry.core.constraints.generic.union_constraint import UnionConstraint
 
 
@@ -17,4 +18,5 @@
     "MultiConstraint",
     "UnionConstraint",
     "parse_constraint",
+    "parse_extra_constraint",
 )

src/poetry/core/constraints/generic/constraint.py

@@ -137,7 +137,7 @@ def allows_any(self, other: BaseConstraint) -> bool:
         return other.is_any()
 
     def invert(self) -> Constraint:
-        return Constraint(self._value, self._trans_op_inv[self.operator])
+        return self.__class__(self._value, self._trans_op_inv[self.operator])
 
     def difference(self, other: BaseConstraint) -> Constraint | EmptyConstraint:
         if other.allows(self):
@@ -207,8 +207,8 @@ def is_empty(self) -> bool:
         return False
 
     def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Constraint):
-            return NotImplemented
+        if not isinstance(other, self.__class__):
+            return False
 
         return (self.value, self.operator) == (other.value, other.operator)
 
@@ -220,3 +220,44 @@ def __str__(self) -> str:
             return f"'{self._value}' {self._operator}"
         op = self._operator if self._operator != "==" else ""
         return f"{op}{self._value}"
+
+
+class ExtraConstraint(Constraint):
+    def __init__(self, value: str, operator: str = "==") -> None:
+        super().__init__(value, operator)
+        # Do the check after calling the super constructor,
+        # i.e. after the operator has been normalized.
+        if self._operator not in {"==", "!="}:
+            raise ValueError(
+                'Only the operators "==" and "!=" are supported for extra constraints'
+            )
+
+    def intersect(self, other: BaseConstraint) -> BaseConstraint:
+        from poetry.core.constraints.generic.multi_constraint import (
+            ExtraMultiConstraint,
+        )
+
+        if isinstance(other, Constraint):
+            if other == self:
+                return self
+
+            if self._value == other._value and self._operator != other.operator:
+                return EmptyConstraint()
+
+            return ExtraMultiConstraint(self, other)
+
+        return super().intersect(other)
+
+    def union(self, other: BaseConstraint) -> BaseConstraint:
+        from poetry.core.constraints.generic.union_constraint import UnionConstraint
+
+        if isinstance(other, Constraint):
+            if other == self:
+                return self
+
+            if self._value == other._value and self._operator != other.operator:
+                return AnyConstraint()
+
+            return UnionConstraint(self, other)
+
+        return super().union(other)

src/poetry/core/constraints/generic/multi_constraint.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import itertools
+
 from typing import TYPE_CHECKING
 
 from poetry.core.constraints.generic import AnyConstraint
@@ -13,8 +15,10 @@
 
 
 class MultiConstraint(BaseConstraint):
+    OPERATORS: tuple[str, ...] = ("!=", "in", "not in")
+
     def __init__(self, *constraints: Constraint) -> None:
-        if any(c.operator == "==" for c in constraints):
+        if any(c.operator not in self.OPERATORS for c in constraints):
             raise ValueError(
                 "A multi-constraint can only be comprised of negative constraints"
             )
@@ -62,7 +66,7 @@ def intersect(self, other: BaseConstraint) -> BaseConstraint:
             union = list(self.constraints) + [
                 c for c in other.constraints if c not in ours
             ]
-            return MultiConstraint(*union)
+            return self.__class__(*union)
 
         if not isinstance(other, Constraint):
             return other.intersect(self)
@@ -74,16 +78,16 @@ def intersect(self, other: BaseConstraint) -> BaseConstraint:
             # same value but different operator, e.g. '== "linux"' and '!= "linux"'
             return EmptyConstraint()
 
-        if other.operator == "==":
+        if other.operator == "==" and "==" not in self.OPERATORS:
             return other
 
-        return MultiConstraint(*self._constraints, other)
+        return self.__class__(*self._constraints, other)
 
     def union(self, other: BaseConstraint) -> BaseConstraint:
         if isinstance(other, MultiConstraint):
             theirs = set(other.constraints)
             common = [c for c in self.constraints if c in theirs]
-            return MultiConstraint(*common)
+            return self.__class__(*common)
 
         if not isinstance(other, Constraint):
             return other.union(self)
@@ -102,10 +106,10 @@ def union(self, other: BaseConstraint) -> BaseConstraint:
         if len(constraints) == 1:
             return constraints[0]
 
-        return MultiConstraint(*constraints)
+        return self.__class__(*constraints)
 
     def __eq__(self, other: object) -> bool:
-        if not isinstance(other, MultiConstraint):
+        if not isinstance(other, self.__class__):
             return False
 
         return self._constraints == other._constraints
@@ -116,3 +120,50 @@ def __hash__(self) -> int:
     def __str__(self) -> str:
         constraints = [str(constraint) for constraint in self._constraints]
         return ", ".join(constraints)
+
+
+class ExtraMultiConstraint(MultiConstraint):
+    # Since the extra marker can have multiple values at the same time,
+    # "==extra1, ==extra2" is not empty!
+    OPERATORS = ("==", "!=")
+
+    def intersect(self, other: BaseConstraint) -> BaseConstraint:
+        if isinstance(other, MultiConstraint):
+            op_values = {}
+            for op in self.OPERATORS:
+                op_values[op] = {
+                    c.value
+                    for c in itertools.chain(self._constraints, other.constraints)
+                    if c.operator == op
+                }
+            if op_values["=="] & op_values["!="]:
+                return EmptyConstraint()
+
+        return super().intersect(other)
+
+    def union(self, other: BaseConstraint) -> BaseConstraint:
+        from poetry.core.constraints.generic import UnionConstraint
+
+        if isinstance(other, MultiConstraint):
+            if set(other.constraints) == set(self._constraints):
+                return self
+            return UnionConstraint(self, other)
+
+        if isinstance(other, Constraint):
+            if other in self._constraints:
+                return other
+
+            if len(self._constraints) == 2 and other.value in (
+                c.value for c in self._constraints
+            ):
+                # same value but different operator
+                constraints: list[BaseConstraint] = [
+                    *(c for c in self._constraints if c.value != other.value),
+                    other,
+                ]
+            else:
+                constraints = [self, other]
+
+            return UnionConstraint(*constraints)
+
+        return super().union(other)

src/poetry/core/constraints/generic/parser.py

@@ -7,6 +7,7 @@
 
 from poetry.core.constraints.generic.any_constraint import AnyConstraint
 from poetry.core.constraints.generic.constraint import Constraint
+from poetry.core.constraints.generic.constraint import ExtraConstraint
 from poetry.core.constraints.generic.union_constraint import UnionConstraint
 from poetry.core.constraints.version.exceptions import ParseConstraintError
 
@@ -29,6 +30,17 @@
 
 @functools.cache
 def parse_constraint(constraints: str) -> BaseConstraint:
+    return _parse_constraint(constraints, Constraint)
+
+
+@functools.cache
+def parse_extra_constraint(constraints: str) -> BaseConstraint:
+    return _parse_constraint(constraints, ExtraConstraint)
+
+
+def _parse_constraint(
+    constraints: str, constraint_type: type[Constraint]
+) -> BaseConstraint:
     if constraints == "*":
         return AnyConstraint()
 
@@ -40,9 +52,13 @@ def parse_constraint(constraints: str) -> BaseConstraint:
 
         if len(and_constraints) > 1:
             for constraint in and_constraints:
-                constraint_objects.append(parse_single_constraint(constraint))
+                constraint_objects.append(
+                    _parse_single_constraint(constraint, constraint_type)
+                )
         else:
-            constraint_objects.append(parse_single_constraint(and_constraints[0]))
+            constraint_objects.append(
+                _parse_single_constraint(and_constraints[0], constraint_type)
+            )
 
         if len(constraint_objects) == 1:
             constraint = constraint_objects[0]
@@ -59,12 +75,14 @@ def parse_constraint(constraints: str) -> BaseConstraint:
         return UnionConstraint(*or_groups)
 
 
-def parse_single_constraint(constraint: str) -> Constraint:
+def _parse_single_constraint(
+    constraint: str, constraint_type: type[Constraint]
+) -> Constraint:
     # string comparator
     if m := STR_CMP_CONSTRAINT.match(constraint):
         op = m.group("op")
         value = m.group("value").strip()
-        return Constraint(value, op)
+        return constraint_type(value, op)
 
     # Basic comparator
 
@@ -75,6 +93,6 @@ def parse_single_constraint(constraint: str) -> Constraint:
 
         version = m.group(2).strip()
 
-        return Constraint(version, op)
+        return constraint_type(version, op)
 
     raise ParseConstraintError(f"Could not parse version constraint: {constraint}")

src/poetry/core/constraints/generic/union_constraint.py

@@ -5,7 +5,9 @@
 from poetry.core.constraints.generic import AnyConstraint
 from poetry.core.constraints.generic.base_constraint import BaseConstraint
 from poetry.core.constraints.generic.constraint import Constraint
+from poetry.core.constraints.generic.constraint import ExtraConstraint
 from poetry.core.constraints.generic.empty_constraint import EmptyConstraint
+from poetry.core.constraints.generic.multi_constraint import ExtraMultiConstraint
 from poetry.core.constraints.generic.multi_constraint import MultiConstraint
 
 
@@ -48,7 +50,11 @@ def invert(self) -> MultiConstraint:
             raise NotImplementedError(
                 "Inversion of complex union constraints not implemented"
             )
-        return MultiConstraint(*inverted_constraints)  # type: ignore[arg-type]
+        if any(isinstance(c, ExtraConstraint) for c in inverted_constraints):
+            multi_type: type[MultiConstraint] = ExtraMultiConstraint
+        else:
+            multi_type = MultiConstraint
+        return multi_type(*inverted_constraints)  # type: ignore[arg-type]
 
     def intersect(self, other: BaseConstraint) -> BaseConstraint:
         if other.is_any():
@@ -57,6 +63,14 @@ def intersect(self, other: BaseConstraint) -> BaseConstraint:
         if other.is_empty():
             return other
 
+        if isinstance(other, UnionConstraint) and set(other.constraints) == set(
+            self._constraints
+        ):
+            return self
+
+        if isinstance(other, ExtraConstraint) and other in self._constraints:
+            return other
+
         if isinstance(other, Constraint):
             # (A or B) and C => (A and C) or (B and C)
             # just a special case of UnionConstraint
@@ -99,6 +113,9 @@ def union(self, other: BaseConstraint) -> BaseConstraint:
         if other.is_empty():
             return self
 
+        if other == self:
+            return self
+
         if isinstance(other, Constraint):
             # (A or B) or C => A or B or C
             # just a special case of UnionConstraint