Union with subclasses

[csqzhang]

• cattrs version: 24.1.2
• Python version: 3.11
• Operating System: Linux/MacOSX

Description

I am not sure how to handle union with subclass.

What I Did

import attr
import cattr
from cattrs.strategies import configure_tagged_union, include_subclasses


@attr.s(frozen=True, auto_attribs=True)
class A:
    pass


@attr.s(frozen=True, auto_attribs=True)
class B:
    pass


@attr.s(frozen=True, auto_attribs=True)
class B2(B):
    b: int



@attr.s(frozen=True, auto_attribs=True)
class C:
    x: A | B


CONVERTER = cattr.Converter()
configure_tagged_union(A | B, CONVERTER)
union_strategy = partial(configure_tagged_union, tag_name="type_name")
include_subclasses(B, CONVERTER, union_strategy=union_strategy)

if __name__ == "__main__":
    instance = C(x=B2(b=1))
    r = CONVERTER.unstructure(instance, unstructure_as=C)
    print(r)
    print(CONVERTER.structure(r, C))

Getting error as

Traceback (most recent call last):
  File "..../IPython/core/interactiveshell.py", line 3577, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-1-8e14e71c7dee>", line 1, in <module>
    CONVERTER.unstructure(instance, unstructure_as=C)
  File ".../python3.11/site-packages/cattrs/converters.py", line 300, in unstructure
    return self._unstructure_func.dispatch(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<cattrs generated unstructure __main__.C>", line 3, in unstructure_C
    'x': __c_unstr_x(instance.x),
          ^^^^^^^^^^^^^^^^^^^^^^^
  File ".../python3.11/site-packages/cattrs/strategies/_unions.py", line 82, in unstructure_tagged_union
    res = _exact_cl_unstruct_hooks[val.__class__](val)
          ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
KeyError: <class '__main__.B2'>
^CTraceback (most recent call last):
  File ".../example.py", line 99, in <module>
    r = CONVERTER.unstruct

My understanding is that configure_tagged_union does not register all the hooks for subclasses. I wonder what is the right/best way to workaround it?

[Tinche]

Ah, I see the issue. configure_tagged_union doesn't compose well with include_subclasses (the reverse direction works ok). You're trying to do two layers - the first layer is A | B, and the second is all the subclasses of B.

We can work around by overriding part of configure_tagged_union ourselves (the unstructure hook). That strategy is pretty simple so it won't be a lot of work.

Here's the code:

from functools import partial

from attrs import frozen

import cattr
from cattrs.strategies import configure_tagged_union, include_subclasses


@frozen
class A:
    pass


@frozen
class B:
    pass


@frozen
class B2(B):
    b: int


@frozen
class C:
    x: A | B


CONVERTER = cattr.Converter()

union_strategy = partial(configure_tagged_union, tag_name="type_name")
include_subclasses(B, CONVERTER, union_strategy=union_strategy)

configure_tagged_union(A | B, CONVERTER)


@CONVERTER.register_unstructure_hook
def unstructure_a_b(val: A | B) -> dict:
    res = CONVERTER.unstructure(val)
    res["_type"] = "A" if isinstance(val, A) else "B"
    return res


if __name__ == "__main__":
    instance = C(x=B2(b=1))
    r = CONVERTER.unstructure(instance, unstructure_as=C)
    print(r)
    print(CONVERTER.structure(r, C))

I also flipped the order in which strategies are applied (it matters).

The unstructured payload will have two keys for the class, type_name and _type. We could probably reduce that to one by doing some more overriding. Let me know if this is interesting!