Tweak the include_subclasses strategy

Author: Tinche

src/cattrs/strategies/_subclasses.py

@@ -1,9 +1,14 @@
 """Strategies for customizing subclass behaviors."""
 from gc import collect
-from typing import Dict, Optional, Tuple, Type, Union, List, Callable, Any, get_args
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, get_args
 
-from ..converters import Converter, BaseConverter
-from ..gen import AttributeOverride, make_dict_structure_fn, make_dict_unstructure_fn
+from ..converters import BaseConverter, Converter
+from ..gen import (
+    AttributeOverride,
+    _already_generating,
+    make_dict_structure_fn,
+    make_dict_unstructure_fn,
+)
 
 
 def _make_subclasses_tree(cl: Type) -> List[Type]:
@@ -12,7 +17,8 @@ def _make_subclasses_tree(cl: Type) -> List[Type]:
     ]
 
 
-def _has_subclasses(cl: Type, given_subclasses: Tuple[Type]):
+def _has_subclasses(cl: Type, given_subclasses: Tuple[Type, ...]) -> bool:
+    """Whether the given class has subclasses from `given_subclasses`."""
     actual = set(cl.__subclasses__())
     given = set(given_subclasses)
     return bool(actual & given)
@@ -31,7 +37,7 @@ def _get_union_type(cl: Type, given_subclasses_tree: Tuple[Type]) -> Optional[Ty
 def include_subclasses(
     cl: Type,
     converter: Converter,
-    subclasses: Optional[Tuple[Type]] = None,
+    subclasses: Optional[Tuple[Type, ...]] = None,
     union_strategy: Optional[Callable[[Any, BaseConverter], Any]] = None,
     overrides: Optional[Dict[str, AttributeOverride]] = None,
 ) -> None:
@@ -139,40 +145,76 @@ def unstruct_hook(
 
 def _include_subclasses_with_union_strategy(
     converter: Converter,
-    union_classes: Tuple[Type],
+    union_classes: Tuple[Type, ...],
     union_strategy: Callable[[Any, BaseConverter], Any],
     overrides: Dict[str, AttributeOverride],
 ):
+    """
+    This function is tricky because we're dealing with what is essentially a circular reference.
+
+    We need to generate a structure hook for a class that is both:
+    * specific for that particular class and its own fields
+    * but should handle specific functions for all its descendants too
+
+    Hence the dance with registering below.
+    """
+
     parent_classes = [cl for cl in union_classes if _has_subclasses(cl, union_classes)]
     if not parent_classes:
         return
 
+    original_unstruct_hooks = {}
+    original_struct_hooks = {}
     for cl in union_classes:
+        # In the first pass, every class gets its own unstructure function according to
+        # the overrides.
+        # We just generate the hooks, and do not register them. This allows us to manipulate
+        # the _already_generating set to force runtime dispatch.
+        _already_generating.working_set = set(union_classes) - {cl}
+        try:
+            unstruct_hook = make_dict_unstructure_fn(cl, converter, **overrides)
+            struct_hook = make_dict_structure_fn(cl, converter, **overrides)
+        finally:
+            _already_generating.working_set = set()
+        original_unstruct_hooks[cl] = unstruct_hook
+        original_struct_hooks[cl] = struct_hook
+
+    # Now that's done, we can register all the hooks and generate the
+    # union handler. The union handler needs them.
+    final_union = Union[union_classes]  # type: ignore
+
+    for cl, hook in original_unstruct_hooks.items():
 
         def cls_is_cl(cls, _cl=cl):
             return cls is _cl
 
-        converter.register_structure_hook_func(
-            cls_is_cl, make_dict_structure_fn(cl, converter, **overrides)
-        )
-        converter.register_unstructure_hook_func(
-            cls_is_cl, make_dict_unstructure_fn(cl, converter, **overrides)
-        )
+        converter.register_unstructure_hook_func(cls_is_cl, hook)
 
-    for cl in parent_classes:
-        subclass_union = _get_union_type(cl, union_classes)
-        sub_union_classes = get_args(subclass_union)
-        union_strategy(subclass_union, converter)
-        struct_hook = converter._union_struct_registry[subclass_union]
-        unstruct_hook = converter._unstructure_func.dispatch(subclass_union)
+    for cl, hook in original_struct_hooks.items():
 
         def cls_is_cl(cls, _cl=cl):
             return cls is _cl
 
-        def cls_is_in_union(cls, _union_classes=sub_union_classes):
-            return cls in _union_classes
+        converter.register_structure_hook_func(cls_is_cl, hook)
 
-        # This needs to use function dispatch, using singledispatch will again
-        # match A and all subclasses, which is not what we want.
-        converter.register_structure_hook_func(cls_is_cl, struct_hook)
-        converter.register_unstructure_hook_func(cls_is_in_union, unstruct_hook)
+    union_strategy(final_union, converter)
+    unstruct_hook = converter._unstructure_func.dispatch(final_union)
+    struct_hook = converter._structure_func.dispatch(final_union)
+
+    for cl in union_classes:
+        # In the second pass, we overwrite the hooks with the union hook.
+
+        def cls_is_cl(cls, _cl=cl):
+            return cls is _cl
+
+        converter.register_unstructure_hook_func(cls_is_cl, unstruct_hook)
+        subclasses = tuple([c for c in union_classes if issubclass(c, cl)])
+        if len(subclasses) > 1:
+            u = Union[subclasses]  # type: ignore
+            union_strategy(u, converter)
+            struct_hook = converter._structure_func.dispatch(u)
+
+            def sh(payload: dict, _, _u=u, _s=struct_hook) -> cl:
+                return _s(payload, _u)
+
+            converter.register_structure_hook_func(cls_is_cl, sh)

tests/strategies/test_include_subclasses.py

@@ -1,13 +1,13 @@
 import typing
-from functools import partial
 from copy import deepcopy
+from functools import partial
+from typing import Tuple
 
 import attr
 import pytest
 
 from cattrs import Converter, override
 from cattrs.errors import ClassValidationError
-from cattrs.strategies._subclasses import _make_subclasses_tree
 from cattrs.strategies import configure_tagged_union, include_subclasses
 
 
@@ -148,8 +148,8 @@ def conv_w_subclasses(request):
     "struct_unstruct", IDS_TO_STRUCT_UNSTRUCT.values(), ids=IDS_TO_STRUCT_UNSTRUCT
 )
 def test_structuring_with_inheritance(
-    conv_w_subclasses: typing.Tuple[Converter, bool], struct_unstruct
-):
+    conv_w_subclasses: Tuple[Converter, bool], struct_unstruct
+) -> None:
     structured, unstructured = struct_unstruct
 
     converter, included_subclasses_param = conv_w_subclasses
@@ -176,7 +176,7 @@ def test_structuring_with_inheritance(
 
     if structured.__class__ in {Parent, Child1, Child2}:
         with pytest.raises(ClassValidationError):
-            converter.structure(unstructured, GrandChild)
+            _ = converter.structure(unstructured, GrandChild)
 
 
 def test_structure_as_union():
@@ -201,14 +201,11 @@ def test_circular_reference(conv_w_subclasses):
     if included_subclasses_param != "with-subclasses-and-tagged-union":
         unstruct = _remove_type_name(unstruct)
 
-    if included_subclasses_param == "wo-subclasses":
-        # We already now that it will fail
-        return
+    if "wo-subclasses" in included_subclasses_param:
+        pytest.xfail("Cannot succeed if include_subclasses strategy is not used")
 
     res = c.unstructure(struct)
-    if "wo-subclasses" or "tagged-union" in included_subclasses_param:
-        # TODO: tagged-union should work here, but it does not yet.
-        pytest.xfail("Cannot succeed if include_subclasses strategy is not used")
+
     assert res == unstruct
 
     res = c.unstructure(struct, CircularA)
@@ -222,7 +219,7 @@ def test_circular_reference(conv_w_subclasses):
     "struct_unstruct", IDS_TO_STRUCT_UNSTRUCT.values(), ids=IDS_TO_STRUCT_UNSTRUCT
 )
 def test_unstructuring_with_inheritance(
-    conv_w_subclasses: typing.Tuple[Converter, bool], struct_unstruct
+    conv_w_subclasses: Tuple[Converter, bool], struct_unstruct
 ):
     structured, unstructured = struct_unstruct
     converter, included_subclasses_param = conv_w_subclasses
@@ -326,29 +323,3 @@ def test_overrides(with_union_strategy: bool, struct_unstruct: str):
     assert c.unstructure(structured) == unstructured
     assert c.structure(unstructured, Parent) == structured
     assert c.structure(unstructured, structured.__class__) == structured
-
-
-def test_class_tree_generator():
-    class P:
-        pass
-
-    class C1(P):
-        pass
-
-    class C2(P):
-        pass
-
-    class GC1(C2):
-        pass
-
-    class GC2(C2):
-        pass
-
-    tree_c1 = _make_subclasses_tree(C1)
-    assert tree_c1 == [C1]
-
-    tree_c2 = _make_subclasses_tree(C2)
-    assert tree_c2 == [C2, GC1, GC2]
-
-    tree_p = _make_subclasses_tree(P)
-    assert tree_p == [P, C1, C2, GC1, GC2]