Lupl/mutual exclusion mixin

tests/tests_pydantic_tools/test_mutual_constraint_validator.py

@@ -0,0 +1,133 @@
+"""Pytest entry point for mutual constraint validation mixin."""
+
+from typing import NamedTuple
+
+import pytest
+
+from pydantic import BaseModel, ConfigDict
+from upto import _MutualConstraintMixin
+from upto.pydantic_tools.mutual_constraint_validator import (
+    MutualDependencyException,
+    MutualExclusionException,
+)
+
+
+class Model(BaseModel, _MutualConstraintMixin):
+    a: int | None = None
+    b: int | None = None
+    c: int | None = None
+    d: int | None = None
+
+
+class MutualConstraintParameter(NamedTuple):
+    model: BaseModel
+    model_config: ConfigDict
+    model_init: dict
+    exception: Exception | None = None
+
+
+mutual_constraint_pass_parameters = [
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(),
+        model_init={"a": 1, "c": 3},
+        exception=None,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_exclusive=[("a", "b")]),
+        model_init={"a": 1, "c": 3},
+        exception=None,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_exclusive=[("a", "b")]),
+        model_init={"b": 2, "c": 3, "d": 4},
+        exception=None,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_exclusive=[("a", "d")]),
+        model_init={"a": 1, "b": 2, "c": 3},
+        exception=None,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_dependent=[("a", "b")]),
+        model_init={"a": 1, "b": 2, "c": 3},
+        exception=MutualDependencyException,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_dependent=[("a", "b")]),
+        model_init={"c": 3, "d": 4, "a": 1, "b": 2},
+        exception=MutualDependencyException,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_dependent=[("a", "d")]),
+        model_init={"a": 1, "c": 3, "d": 4},
+        exception=MutualDependencyException,
+    ),
+]
+
+mutual_constraint_fail_parameters: list[MutualConstraintParameter] = [
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_exclusive=[("a", "b")]),
+        model_init={"a": 1, "b": 2},
+        exception=MutualExclusionException,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_exclusive=[("a", "b")]),
+        model_init={"a": 1, "b": 2, "c": 3, "d": 4},
+        exception=MutualExclusionException,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_exclusive=[("a", "d")]),
+        model_init={"a": 1, "b": 2, "c": 3, "d": 4},
+        exception=MutualExclusionException,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_dependent=[("a", "b")]),
+        model_init={"a": 1, "c": 3},
+        exception=MutualDependencyException,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_dependent=[("a", "b")]),
+        model_init={"c": 3, "d": 4},
+        exception=MutualDependencyException,
+    ),
+    MutualConstraintParameter(
+        model=Model,
+        model_config=ConfigDict(mutually_dependent=[("a", "d")]),
+        model_init={"c": 3, "d": 4},
+        exception=MutualDependencyException,
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    MutualConstraintParameter._fields, mutual_constraint_pass_parameters
+)
+def test_happy_path_mutual_constrain_validator(
+    model, model_config, model_init, exception
+):
+    model.model_config = model_config
+    assert model(**model_init)
+
+
+@pytest.mark.parametrize(
+    MutualConstraintParameter._fields, mutual_constraint_fail_parameters
+)
+def test_sad_path_mutual_constrain_validator(
+    model, model_config, model_init, exception
+):
+    model.model_config = model_config
+
+    with pytest.raises(exception):
+        model(**model_init)

upto/__init__.py

@@ -1,3 +1,4 @@
 from upto.compose_router import ComposeRouter
 from upto.pydantic_tools.curry_model import CurryModel
 from upto.pydantic_tools.model_constructors import init_model_from_kwargs
+from upto.pydantic_tools.mutual_constraint_validator import _MutualConstraintMixin

upto/pydantic_tools/mutual_constraint_validator.py

@@ -0,0 +1,78 @@
+"""Mixin enabling Pydantic models to check for mutually dependent fields."""
+
+from typing import NoReturn, Protocol
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+
+class MutualExclusionException(Exception):
+    pass
+
+
+class MutualDependencyException(Exception):
+    pass
+
+
+class _MutualConstraintTemplate(Protocol):
+    constraint_id: str
+
+    def condition(self, pair: tuple[str, str], model_data: dict) -> bool: ...
+
+    def throw(self, pair: tuple[str, str]) -> NoReturn: ...
+
+
+class MutualExclusionTemplate(_MutualConstraintTemplate):
+    constraint_id = "mutually_exclusive"
+
+    def condition(self, pair: tuple[str, str], model_data: dict) -> bool:
+        return all(field in model_data for field in pair)
+
+    def throw(self, pair: tuple[str, str]) -> NoReturn:
+        raise MutualExclusionException(
+            f"Fields {' and '.join(pair)} are mutually exclusive."
+        )
+
+
+class MutualDependencyTemplate(_MutualConstraintTemplate):
+    constraint_id = "mutually_dependent"
+
+    def condition(self, pair: tuple[str, str], model_data: dict) -> bool:
+        return not all(field in model_data for field in pair)
+
+    def throw(self, pair: tuple[str, str]) -> NoReturn:
+        raise MutualDependencyException(
+            f"Fields {' and '.join(pair)} are mutually dependent."
+        )
+
+
+class _MutualConstraintMixin:
+    """Mixin for adding mutual exclusion and dependency validators to a model.
+
+    Subclass _MutualConstrainMixin in a Pydantic model definition
+    to add the validator to the model.
+
+    The validator reads two keys from model_config:
+
+    - mutually_exclusive: Iterble[tuple[str, str]]
+    - mutually_dependent: Iterble[tuple[str, str]]
+
+    and checks the model fields accordingly.
+    """
+
+    constraints = [
+        MutualExclusionTemplate(),
+        MutualDependencyTemplate(),
+    ]
+
+    @model_validator(mode="before")
+    @classmethod
+    def _check_mutual_constraints(cls, data: dict) -> dict:
+        config = cls.model_config  # type: ignore
+
+        for constraint in cls.constraints:
+            if pairs := config.get(constraint.constraint_id):
+                for pair in pairs:
+                    if constraint.condition(pair, data):
+                        constraint.throw(pair)
+
+        return data