Implement a disambiguator for discriminated unions (#392)

* Implement a disambiguator for discriminated unions

Another way at "tagged unions" except backwards compatible! Also, a ton
of services provide data in this fashion. Debatably cattrs already does
something similar where you don't need to register a structure hook for
a `Literal[1, 2, 3, 4]` or whatever.

* PR feedback

* Stop supporting Python 3.7

* Remove forgotten `is_py37` reference...

* Fix ruff error + one isort error

* Don't add `isort` to CI

* Oops, don't fix ruff and then error out

Author: A5rocks

HISTORY.md

@@ -47,6 +47,8 @@
   ([#393](https://github.com/python-attrs/cattrs/issues/393))
 - Remove some unused lines in the unstructuring code.
   ([#416](https://github.com/python-attrs/cattrs/pull/416))
+- Disambiguate a union of attrs classes where there's a `typing.Literal` tag of some sort.
+  ([#391](https://github.com/python-attrs/cattrs/pull/391))
 
 ## 23.1.2 (2023-06-02)
 

Makefile

@@ -47,7 +47,7 @@ clean-test: ## remove test and coverage artifacts
 	rm -f .coverage
 	rm -fr htmlcov/
 
-lint: ## check style with flake8
+lint: ## check style with ruff and black
 	pdm run ruff src/ tests
 	pdm run black --check src tests docs/conf.py
 

src/cattrs/converters.py

@@ -55,7 +55,7 @@
     is_typeddict,
     is_union_type,
 )
-from .disambiguators import create_uniq_field_dis_func
+from .disambiguators import create_default_dis_func
 from .dispatch import MultiStrategyDispatch
 from .errors import (
     IterableValidationError,
@@ -729,13 +729,16 @@ def _get_dis_func(union: Any) -> Callable[[Any], Type]:
                 e for e in union_types if e is not NoneType  # type: ignore
             )
 
+        # TODO: technically both disambiguators could support TypedDicts and
+        # dataclasses...
         if not all(has(get_origin(e) or e) for e in union_types):
             raise StructureHandlerNotFoundError(
                 "Only unions of attrs classes supported "
                 "currently. Register a loads hook manually.",
                 type_=union,
             )
-        return create_uniq_field_dis_func(*union_types)
+
+        return create_default_dis_func(*union_types)
 
     def __deepcopy__(self, _) -> "BaseConverter":
         return self.copy()

src/cattrs/disambiguators.py

@@ -1,22 +1,77 @@
 """Utilities for union (sum type) disambiguation."""
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from functools import reduce
 from operator import or_
-from typing import Any, Callable, Dict, Mapping, Optional, Type
+from typing import Any, Callable, Dict, Mapping, Optional, Type, Union
 
-from attr import NOTHING, fields
+from attr import NOTHING, fields, fields_dict
 
-from cattrs._compat import get_origin
+from cattrs._compat import get_args, get_origin, is_literal
 
 
-def create_uniq_field_dis_func(
+def create_default_dis_func(
     *classes: Type[Any],
 ) -> Callable[[Mapping[Any, Any]], Optional[Type[Any]]]:
     """Given attr classes, generate a disambiguation function.
 
-    The function is based on unique fields."""
+    The function is based on unique fields or unique values."""
     if len(classes) < 2:
         raise ValueError("At least two classes required.")
+
+    # first, attempt for unique values
+
+    # requirements for a discriminator field:
+    # (... TODO: a single fallback is OK)
+    #  - it must be *required*
+    #  - it must always be enumerated
+    cls_candidates = [
+        {
+            at.name
+            for at in fields(get_origin(cl) or cl)
+            if at.default is NOTHING and is_literal(at.type)
+        }
+        for cl in classes
+    ]
+
+    discriminators = cls_candidates[0]
+    for possible_discriminators in cls_candidates:
+        discriminators &= possible_discriminators
+
+    best_result = None
+    best_discriminator = None
+    for discriminator in discriminators:
+        mapping = defaultdict(list)
+
+        for cl in classes:
+            for key in get_args(fields_dict(get_origin(cl) or cl)[discriminator].type):
+                mapping[key].append(cl)
+
+        if best_result is None or max(len(v) for v in mapping.values()) <= max(
+            len(v) for v in best_result.values()
+        ):
+            best_result = mapping
+            best_discriminator = discriminator
+
+    if (
+        best_result
+        and best_discriminator
+        and max(len(v) for v in best_result.values()) != len(classes)
+    ):
+        final_mapping = {
+            k: v[0] if len(v) == 1 else Union[tuple(v)] for k, v in best_result.items()
+        }
+
+        def dis_func(data: Mapping[Any, Any]) -> Optional[Type]:
+            if not isinstance(data, Mapping):
+                raise ValueError("Only input mappings are supported.")
+            return final_mapping[data[best_discriminator]]
+
+        return dis_func
+
+    # next, attempt for unique keys
+
+    # NOTE: This could just as well work with just field availability and not
+    #  uniqueness, returning Unions ... it doesn't do that right now.
     cls_and_attrs = [
         (cl, {at.name for at in fields(get_origin(cl) or cl)}) for cl in classes
     ]
@@ -57,3 +112,6 @@ def dis_func(data: Mapping[Any, Any]) -> Optional[Type]:
         return fallback
 
     return dis_func
+
+
+create_uniq_field_dis_func = create_default_dis_func

tests/test_disambigutors.py

@@ -1,12 +1,12 @@
 """Tests for auto-disambiguators."""
-from typing import Any
+from typing import Any, Literal, Union
 
 import attr
 import pytest
-from attr import NOTHING
+from attrs import NOTHING, define
 from hypothesis import HealthCheck, assume, given, settings
 
-from cattrs.disambiguators import create_uniq_field_dis_func
+from cattrs.disambiguators import create_default_dis_func, create_uniq_field_dis_func
 
 from .untyped import simple_classes
 
@@ -22,6 +22,9 @@ class A:
         # Can't generate for only one class.
         create_uniq_field_dis_func(A)
 
+    with pytest.raises(ValueError):
+        create_default_dis_func(A)
+
     @attr.s
     class B:
         pass
@@ -30,6 +33,9 @@ class B:
         # No fields on either class.
         create_uniq_field_dis_func(A, B)
 
+    with pytest.raises(ValueError):
+        create_default_dis_func(A, B)
+
     @attr.s
     class C:
         a = attr.ib()
@@ -42,6 +48,10 @@ class D:
         # No unique fields on either class.
         create_uniq_field_dis_func(C, D)
 
+    with pytest.raises(ValueError):
+        # No discriminator candidates
+        create_default_dis_func(C, D)
+
     @attr.s
     class E:
         pass
@@ -54,6 +64,18 @@ class F:
         # no usable non-default attributes
         create_uniq_field_dis_func(E, F)
 
+    @define()
+    class G:
+        x: Literal[1]
+
+    @define()
+    class H:
+        x: Literal[1]
+
+    with pytest.raises(ValueError):
+        # The discriminator chosen does not actually help
+        create_default_dis_func(C, D)
+
 
 @given(simple_classes(defaults=False))
 def test_fallback(cl_and_vals):
@@ -99,3 +121,71 @@ def test_disambiguation(cl_and_vals_a, cl_and_vals_b):
     fn = create_uniq_field_dis_func(cl_a, cl_b)
 
     assert fn(attr.asdict(cl_a(*vals_a, **kwargs_a))) is cl_a
+
+
+# not too sure of properties of `create_default_dis_func`
+def test_disambiguate_from_discriminated_enum():
+    # can it find any discriminator?
+    @define()
+    class A:
+        a: Literal[0]
+
+    @define()
+    class B:
+        a: Literal[1]
+
+    fn = create_default_dis_func(A, B)
+    assert fn({"a": 0}) is A
+    assert fn({"a": 1}) is B
+
+    # can it find the better discriminator?
+    @define()
+    class C:
+        a: Literal[0]
+        b: Literal[1]
+
+    @define()
+    class D:
+        a: Literal[0]
+        b: Literal[0]
+
+    fn = create_default_dis_func(C, D)
+    assert fn({"a": 0, "b": 1}) is C
+    assert fn({"a": 0, "b": 0}) is D
+
+    # can it handle multiple tiers of discriminators?
+    # (example inspired by Discord's gateway's discriminated union)
+    @define()
+    class E:
+        op: Literal[1]
+
+    @define()
+    class F:
+        op: Literal[0]
+        t: Literal["MESSAGE_CREATE"]
+
+    @define()
+    class G:
+        op: Literal[0]
+        t: Literal["MESSAGE_UPDATE"]
+
+    fn = create_default_dis_func(E, F, G)
+    assert fn({"op": 1}) is E
+    assert fn({"op": 0, "t": "MESSAGE_CREATE"}) is Union[F, G]
+
+    # can it handle multiple literals?
+    @define()
+    class H:
+        a: Literal[1]
+
+    @define()
+    class J:
+        a: Literal[0, 1]
+
+    @define()
+    class K:
+        a: Literal[0]
+
+    fn = create_default_dis_func(H, J, K)
+    assert fn({"a": 1}) is Union[H, J]
+    assert fn({"a": 0}) is Union[J, K]

tests/test_structure_attrs.py

@@ -1,7 +1,7 @@
 """Loading of attrs classes."""
 from enum import Enum
 from ipaddress import IPv4Address, IPv6Address, ip_address
-from typing import Union
+from typing import Literal, Union
 from unittest.mock import Mock
 
 import pytest
@@ -139,8 +139,6 @@ def dis(obj, _):
 @pytest.mark.parametrize("converter_cls", [BaseConverter, Converter])
 def test_structure_literal(converter_cls):
     """Structuring a class with a literal field works."""
-    from typing import Literal
-
     converter = converter_cls()
 
     @define
@@ -155,8 +153,6 @@ class ClassWithLiteral:
 @pytest.mark.parametrize("converter_cls", [BaseConverter, Converter])
 def test_structure_literal_enum(converter_cls):
     """Structuring a class with a literal field works."""
-    from typing import Literal
-
     converter = converter_cls()
 
     class Foo(Enum):
@@ -175,8 +171,6 @@ class ClassWithLiteral:
 @pytest.mark.parametrize("converter_cls", [BaseConverter, Converter])
 def test_structure_literal_multiple(converter_cls):
     """Structuring a class with a literal field works."""
-    from typing import Literal
-
     converter = converter_cls()
 
     class Foo(Enum):
@@ -210,8 +204,6 @@ class ClassWithLiteral:
 @pytest.mark.parametrize("converter_cls", [BaseConverter, Converter])
 def test_structure_literal_error(converter_cls):
     """Structuring a class with a literal field can raise an error."""
-    from typing import Literal
-
     converter = converter_cls()
 
     @define
@@ -225,8 +217,6 @@ class ClassWithLiteral:
 @pytest.mark.parametrize("converter_cls", [BaseConverter, Converter])
 def test_structure_literal_multiple_error(converter_cls):
     """Structuring a class with a literal field can raise an error."""
-    from typing import Literal
-
     converter = converter_cls()
 
     @define
@@ -302,3 +292,28 @@ def test_structure_prefers_attrib_converters(converter_type):
 
     attrib_converter.assert_any_call(5)
     assert inst.z == "5"
+
+
+@pytest.mark.parametrize("converter_type", [BaseConverter, Converter])
+def test_structure_multitier_discriminator_union(converter_type):
+    converter = converter_type()
+
+    @define()
+    class E:
+        op: Literal[1]
+
+    @define()
+    class F:
+        op: Literal[0]
+        t: Literal["MESSAGE_CREATE"]
+
+    @define()
+    class G:
+        op: Literal[0]
+        t: Literal["MESSAGE_UPDATE"]
+
+    inst = converter.structure({"op": 1}, Union[E, F, G])
+    assert isinstance(inst, E)
+
+    inst = converter.structure({"op": 0, "t": "MESSAGE_CREATE"}, Union[E, F, G])
+    assert isinstance(inst, F)