Better unions

HISTORY.rst

@@ -2,6 +2,12 @@
 History
 =======
 
+1.2.0 (UNRELEASED)
+------------------
+* ``converter.unstructure`` now supports an optional parameter, `unstructure_as`, which can be used to unstructure something as a different type. Useful for unions.
+* Improve support for union un/structuring hooks. Flesh out docs for advanced union handling.
+  (`#115 <https://github.com/Tinche/cattrs/pull/115>`_)
+
 1.1.2 (2020-11-29)
 ------------------
 * The default disambiguator will not consider non-required fields any more.

docs/customizing.rst

@@ -5,6 +5,14 @@ Customizing class un/structuring
 This section deals with customizing the unstructuring and structuring processes
 in ``cattrs``.
 
+Using ``cattr.gen.GenConverter``
+********************************
+
+The ``cattr.gen`` module contains a ``Converter`` subclass, the ``GenConverter``.
+The ``GenConverter``, upon first encountering an ``attrs`` class, will use
+the generation functions mentioned here to generate the specialized hooks for it,
+register the hooks and use them.
+
 Manual un/structuring hooks
 ***************************
 
@@ -100,10 +108,3 @@ keyword in Python.
     >>> c.structure({'class': 1}, ExampleClass)
     ExampleClass(klass=1)
 
-Using ``cattr.gen.GenConverter``
-********************************
-
-The ``cattr.gen`` module also contains a ``Converter`` subclass, the ``GenConverter``.
-The ``GenConverter``, upon first encountering an ``attrs`` class, will use
-the mentioned generation functions to generate the specialized hooks for it,
-register the hooks and use them.

docs/index.rst

@@ -18,6 +18,7 @@ Contents:
    structuring
    unstructuring
    customizing
+   unions
    contributing
    history
 

docs/unions.rst

@@ -0,0 +1,67 @@
+========================
+Tips for Handling Unions
+========================
+
+This sections contains information for advanced union handling.
+
+As mentioned in the structuring section, ``cattrs`` is able to handle simple
+unions of ``attrs`` classes automatically. More complex cases require
+converter customization (since there are many ways of handling unions).
+
+Unstructuring unions with extra metadata
+****************************************
+
+Let's assume a simple scenario of two classes, ``ClassA`` and ``ClassB`, both
+of which have no distinct fields and so cannot be used automatically with
+``cattrs``.
+
+.. code-block:: python
+
+    @attr.define
+    class ClassA:
+        a_string: str
+
+    @attr.define
+    class ClassB:
+        a_string: str
+
+A naive approach to unstructuring either of these would yield identical
+dictionaries, and not enough information to restructure the classes.
+
+.. code-block:: python
+
+    >>> converter.unstructure(ClassA("test"))
+    {'a_string': 'test'}  # Is this ClassA or ClassB? Who knows!
+
+What we can do is ensure some extra information is present in the
+unstructured data, and then use that information to help structure later.
+
+First, we register an unstructure hook for the `Union[ClassA, ClassB]` type.
+
+.. code-block:: python
+
+    >>> converter.register_unstructure_hook(
+    ...     Union[ClassA, ClassB],
+    ...     lambda o: {"_type": type(o).__name__,  **converter.unstructure(o)}
+    ... )
+    >>> converter.unstructure(ClassA("test"), unstructure_as=Union[ClassA, ClassB])
+    {'_type': 'ClassA', 'a_string': 'test'}
+
+Note that when unstructuring, we had to provide the `unstructure_as` parameter
+or `cattrs` would have just applied the usual unstructuring rules to `ClassA`,
+instead of our special union hook.
+
+Now that the unstructured data contains some information, we can create a
+structuring hook to put it to use:
+
+.. code-block:: python
+
+    >>> converter.register_structure_hook(
+    ...     Union[ClassA, ClassB],
+    ...     lambda o, _: converter.structure(o, ClassA if o["_type"] == "ClassA" else ClassB)
+    ... )
+    >>> converter.structure({"_type": "ClassA", "a_string": "test"}, Union[ClassA, ClassB])
+    ClassA(a_string='test')
+
+In the future, `cattrs` will gain additional tools to make union handling even
+easier and automate generating these hooks.

src/cattr/_compat.py

@@ -1,5 +1,6 @@
 import sys
 from typing import (
+    Any,
     Dict,
     FrozenSet,
     List,
@@ -34,7 +35,7 @@ def get_origin(cl):
     from typing import Union, _GenericAlias
 
     def is_tuple(type):
-        return type is Tuple or (
+        return type in (Tuple, tuple) or (
             type.__class__ is _GenericAlias
             and issubclass(type.__origin__, Tuple)
         )
@@ -46,25 +47,27 @@ def is_union_type(obj):
             and obj.__origin__ is Union
         )
 
-    def is_sequence(type):
-        return type is List or (
+    def is_sequence(type: Any) -> bool:
+        return type in (List, list) or (
             type.__class__ is _GenericAlias
             and type.__origin__ is not Union
             and issubclass(type.__origin__, Sequence)
         )
 
     def is_mutable_set(type):
-        return type.__class__ is _GenericAlias and issubclass(
-            type.__origin__, MutableSet
+        return type is set or (
+            type.__class__ is _GenericAlias
+            and issubclass(type.__origin__, MutableSet)
         )
 
     def is_frozenset(type):
-        return type.__class__ is _GenericAlias and issubclass(
-            type.__origin__, FrozenSet
+        return type is frozenset or (
+            type.__class__ is _GenericAlias
+            and issubclass(type.__origin__, FrozenSet)
         )
 
     def is_mapping(type):
-        return type is Mapping or (
+        return type in (Mapping, dict) or (
             type.__class__ is _GenericAlias
             and issubclass(type.__origin__, Mapping)
         )
@@ -112,7 +115,7 @@ def is_union_type(obj):
             and obj.__origin__ is Union
         )
 
-    def is_sequence(type):
+    def is_sequence(type: Any) -> bool:
         return (
             type in (List, list, Sequence, MutableSequence)
             or (

src/cattr/converters.py

@@ -65,7 +65,7 @@ class Converter(object):
         "_unstructure_attrs",
         "_structure_attrs",
         "_dict_factory",
-        "_union_registry",
+        "_union_struct_registry",
         "_structure_func",
     )
 
@@ -95,12 +95,13 @@ def __init__(
         )
         self._unstructure_func.register_func_list(
             [
-                (_subclass(Mapping), self._unstructure_mapping),
-                (_subclass(Sequence), self._unstructure_seq),
-                (_subclass(Set), self._unstructure_seq),
-                (_subclass(FrozenSet), self._unstructure_seq),
+                (is_mapping, self._unstructure_mapping),
+                (is_sequence, self._unstructure_seq),
+                (is_mutable_set, self._unstructure_seq),
+                (is_frozenset, self._unstructure_seq),
                 (_subclass(Enum), self._unstructure_enum),
                 (_is_attrs_class, self._unstructure_attrs),
+                (is_union_type, self._unstructure_union),
             ]
         )
 
@@ -135,11 +136,15 @@ def __init__(
 
         self._dict_factory = dict_factory
 
-        # Unions are instances now, not classes. We use a different registry.
-        self._union_registry = {}
+        # Unions are instances now, not classes. We use different registries.
+        self._union_struct_registry: Dict[
+            Any, Callable[[Any, Type[T]], T]
+        ] = {}
 
-    def unstructure(self, obj: Any) -> Any:
-        return self._unstructure_func.dispatch(obj.__class__)(obj)
+    def unstructure(self, obj: Any, unstructure_as=None) -> Any:
+        return self._unstructure_func.dispatch(
+            obj.__class__ if unstructure_as is None else unstructure_as
+        )(obj)
 
     @property
     def unstruct_strat(self) -> UnstructureStrategy:
@@ -151,14 +156,19 @@ def unstruct_strat(self) -> UnstructureStrategy:
         )
 
     def register_unstructure_hook(
-        self, cls: Type[T], func: Callable[[T], Any]
+        self, cls: Any, func: Callable[[T], Any]
     ) -> None:
         """Register a class-to-primitive converter function for a class.
 
         The converter function should take an instance of the class and return
         its Python equivalent.
         """
-        self._unstructure_func.register_cls_list([(cls, func)])
+        if is_union_type(cls):
+            self._unstructure_func.register_func_list(
+                [(lambda t: t is cls, func)]
+            )
+        else:
+            self._unstructure_func.register_cls_list([(cls, func)])
 
     def register_unstructure_hook_func(
         self, check_func, func: Callable[[T], Any]
@@ -181,7 +191,7 @@ def register_structure_hook(
         is sometimes needed (for example, when dealing with generic classes).
         """
         if is_union_type(cl):
-            self._union_registry[cl] = func
+            self._union_struct_registry[cl] = func
         else:
             self._structure_func.register_cls_list([(cl, func)])
 
@@ -209,13 +219,19 @@ def unstructure_attrs_asdict(self, obj) -> Dict[str, Any]:
         for a in attrs:
             name = a.name
             v = getattr(obj, name)
-            rv[name] = dispatch(v.__class__)(v)
+            rv[name] = dispatch(a.type or v.__class__)(v)
         return rv
 
     def unstructure_attrs_astuple(self, obj) -> Tuple[Any, ...]:
         """Our version of `attrs.astuple`, so we can call back to us."""
         attrs = obj.__class__.__attrs_attrs__
-        return tuple(self.unstructure(getattr(obj, a.name)) for a in attrs)
+        dispatch = self._unstructure_func.dispatch
+        res = list()
+        for a in attrs:
+            name = a.name
+            v = getattr(obj, name)
+            res.append(dispatch(a.type or v.__class__)(v))
+        return tuple(res)
 
     def _unstructure_enum(self, obj):
         """Convert an enum to its value."""
@@ -242,6 +258,14 @@ def _unstructure_mapping(self, mapping):
             for k, v in mapping.items()
         )
 
+    def _unstructure_union(self, obj):
+        """
+        Unstructure an object as a union.
+
+        By default, just unstructures the instance.
+        """
+        return self._unstructure_func.dispatch(obj.__class__)(obj)
+
     # Python primitives to classes.
 
     def _structure_default(self, obj, cl):
@@ -396,7 +420,7 @@ def _structure_union(self, obj, union):
                 return self._structure_func.dispatch(other)(obj, other)
 
         # Check the union registry first.
-        handler = self._union_registry.get(union)
+        handler = self._union_struct_registry.get(union)
         if handler is not None:
             return handler(obj, union)
 

src/cattr/function_dispatch.py

@@ -1,4 +1,5 @@
 from functools import lru_cache
+from typing import Any, Callable
 
 
 class FunctionDispatch(object):
@@ -14,9 +15,9 @@ class FunctionDispatch(object):
 
     def __init__(self):
         self._handler_pairs = []
-        self.dispatch = lru_cache(64)(self._dispatch)
+        self.dispatch = lru_cache(None)(self._dispatch)
 
-    def register(self, can_handle, func):
+    def register(self, can_handle: Callable[[Any], bool], func):
         self._handler_pairs.insert(0, (can_handle, func))
         self.dispatch.cache_clear()
 

src/cattr/gen.py

@@ -37,6 +37,8 @@ def make_dict_unstructure_fn(cl, converter, omit_if_default=False, **kwargs):
         override = kwargs.pop(attr_name, _neutral)
         kn = attr_name if override.rename is None else override.rename
         d = a.default
+        unstruct_type_name = f"__cattr_type_{attr_name}"
+        globs[unstruct_type_name] = a.type
         if d is not attr.NOTHING and (
             (omit_if_default and override.omit_if_default is not False)
             or override.omit_if_default
@@ -52,18 +54,20 @@ def make_dict_unstructure_fn(cl, converter, omit_if_default=False, **kwargs):
                 else:
                     post_lines.append(f"    if i.{attr_name} != {def_name}():")
                 post_lines.append(
-                    f"        res['{kn}'] = __c_u(i.{attr_name})"
+                    f"        res['{kn}'] = __c_u(i.{attr_name}, unstructure_as={unstruct_type_name})"
                 )
             else:
                 globs[def_name] = d
                 post_lines.append(f"    if i.{attr_name} != {def_name}:")
                 post_lines.append(
-                    f"        res['{kn}'] = __c_u(i.{attr_name})"
+                    f"        res['{kn}'] = __c_u(i.{attr_name}, unstructure_as={unstruct_type_name})"
                 )
 
         else:
             # No default or no override.
-            lines.append(f"        '{kn}': __c_u(i.{attr_name}),")
+            lines.append(
+                f"        '{kn}': __c_u(i.{attr_name}, unstructure_as={unstruct_type_name}),"
+            )
     lines.append("    }")
 
     total_lines = lines + post_lines + ["    return res"]

src/cattr/multistrategy_dispatch.py

@@ -26,7 +26,7 @@ class MultiStrategyDispatch(object):
 
     def __init__(self, fallback_func):
         self._function_dispatch = FunctionDispatch()
-        self._function_dispatch.register(lambda cls: True, fallback_func)
+        self._function_dispatch.register(lambda _: True, fallback_func)
         self._single_dispatch = singledispatch(_DispatchNotFound)
         self.dispatch = lru_cache(maxsize=None)(self._dispatch)
 

tests/metadata/test_genconverter.py

@@ -75,7 +75,7 @@ def test_union_field_roundtrip(cl_and_vals_a, cl_and_vals_b, strat):
     """
     converter = Converter(unstruct_strat=strat)
     cl_a, vals_a = cl_and_vals_a
-    cl_b, vals_b = cl_and_vals_b
+    cl_b, _ = cl_and_vals_b
     a_field_names = {a.name for a in fields(cl_a)}
     b_field_names = {a.name for a in fields(cl_b)}
     assume(a_field_names)
@@ -91,7 +91,10 @@ class C(object):
     inst = C(a=cl_a(*vals_a))
 
     if strat is UnstructureStrategy.AS_DICT:
-        assert inst == converter.structure(converter.unstructure(inst), C)
+        unstructured = converter.unstructure(inst)
+        assert inst == converter.structure(
+            converter.unstructure(unstructured), C
+        )
     else:
         # Our disambiguation functions only support dictionaries for now.
         with pytest.raises(ValueError):
@@ -100,9 +103,9 @@ class C(object):
         def handler(obj, _):
             return converter.structure(obj, cl_a)
 
-        converter._union_registry[Union[cl_a, cl_b]] = handler
-        assert inst == converter.structure(converter.unstructure(inst), C)
-        del converter._union_registry[Union[cl_a, cl_b]]
+        converter.register_structure_hook(Union[cl_a, cl_b], handler)
+        unstructured = converter.unstructure(inst)
+        assert inst == converter.structure(unstructured, C)
 
 
 @given(simple_typed_classes(defaults=False))