Source code for cattrs.strategies._unions

from collections import defaultdict
from typing import Any, Callable, Dict, Optional, Type

from attrs import NOTHING

from cattrs import Converter

__all__ = ["default_tag_generator", "configure_tagged_union"]

def default_tag_generator(typ: Type) -> str:
    """Return the class name."""
    return typ.__name__

[docs]def configure_tagged_union( union: Any, converter: Converter, tag_generator: Callable[[Type], str] = default_tag_generator, tag_name: str = "_type", default: Optional[Type] = NOTHING, ) -> None: """ Configure the converter so that `union` (which should be a union) is un/structured with the help of an additional piece of data in the unstructured payload, the tag. :param converter: The converter to apply the strategy to. :param tag_generator: A `tag_generator` function is used to map each member of the union to a tag, which is then included in the unstructured payload. The default tag generator returns the name of the class. :param tag_name: The key under which the tag will be set in the unstructured payload. By default, `'_type'`. :param default: An optional class to be used if the tag information is not present when structuring. The tagged union strategy currently only works with the dict un/structuring base strategy. .. versionadded:: 23.1.0 """ args = union.__args__ tag_to_hook = {} exact_cl_unstruct_hooks = {} for cl in args: tag = tag_generator(cl) struct_handler = converter._structure_func.dispatch(cl) unstruct_handler = converter._unstructure_func.dispatch(cl) def structure_union_member(val: dict, _cl=cl, _h=struct_handler) -> cl: return _h(val, _cl) def unstructure_union_member(val: union, _h=unstruct_handler) -> dict: return _h(val) tag_to_hook[tag] = structure_union_member exact_cl_unstruct_hooks[cl] = unstructure_union_member cl_to_tag = {cl: tag_generator(cl) for cl in args} if default is not NOTHING: default_handler = converter._structure_func.dispatch(default) def structure_default(val: dict, _cl=default, _h=default_handler): return _h(val, _cl) tag_to_hook = defaultdict(lambda: structure_default, tag_to_hook) cl_to_tag = defaultdict(lambda: default, cl_to_tag) def unstructure_tagged_union( val: union, _exact_cl_unstruct_hooks=exact_cl_unstruct_hooks, _cl_to_tag=cl_to_tag, _tag_name=tag_name, ) -> Dict: res = _exact_cl_unstruct_hooks[val.__class__](val) res[_tag_name] = _cl_to_tag[val.__class__] return res if default is NOTHING: def structure_tagged_union( val: dict, _, _tag_to_cl=tag_to_hook, _tag_name=tag_name ) -> union: return _tag_to_cl[val[_tag_name]](val) else: def structure_tagged_union( val: dict, _, _tag_to_hook=tag_to_hook, _tag_name=tag_name, _dh=default_handler, _default=default, ) -> union: if _tag_name in val: return _tag_to_hook[val[_tag_name]](val) else: return _dh(val, _default) converter.register_unstructure_hook(union, unstructure_tagged_union) converter.register_structure_hook(union, structure_tagged_union)