"""Type-safe, human-readable serialization helpers for dataclasses."""
from __future__ import annotations
import dataclasses
import enum
import functools
from typing import IO, TYPE_CHECKING, Any, Optional, Set, Type, TypeVar, Union
if TYPE_CHECKING:
# Since serialization functionality is deprecated, the yaml dependency is optional.
import yaml
from typing_extensions import get_args, get_origin
from .. import _resolver
from ..constructors._struct_spec import MISSING
ENUM_YAML_TAG_PREFIX = "!enum:"
DATACLASS_YAML_TAG_PREFIX = "!dataclass:"
MISSING_YAML_TAG_PREFIX = "!missing"
DataclassType = TypeVar("DataclassType")
def _get_contained_special_types_from_type(
cls: Type[Any],
_parent_contained_dataclasses: Optional[Set[Type[Any]]] = None,
) -> Set[Type[Any]]:
"""Takes a dataclass type, and recursively searches its fields for dataclass or enum
types."""
assert _resolver.is_dataclass(cls)
parent_contained_dataclasses = (
set()
if _parent_contained_dataclasses is None
else _parent_contained_dataclasses
)
cls = _resolver.unwrap_annotated(cls)
cls, type_from_typevar = _resolver.resolve_generic_types(cls)
contained_special_types = {cls}
def handle_type(typ: Type[Any]) -> Set[Type[Any]]:
# Handle dataclasses.
if _resolver.is_dataclass(typ) and typ not in parent_contained_dataclasses:
return _get_contained_special_types_from_type(
typ,
_parent_contained_dataclasses=contained_special_types
| parent_contained_dataclasses,
)
# Handle enums.
elif isinstance(typ, enum.EnumMeta):
return {typ}
# Handle Union, Annotated, List, etc. No-op when there are no args.
return functools.reduce(set.union, map(handle_type, get_args(typ)), set())
# Handle generics.
for typ in type_from_typevar.values():
contained_special_types |= handle_type(typ)
if cls in parent_contained_dataclasses:
return contained_special_types
# Handle fields.
for field in _resolver.resolved_fields(cls): # type: ignore
assert not isinstance(field.type, str)
contained_special_types |= handle_type(field.type)
# Handle subclasses.
for subclass in cls.__subclasses__():
contained_special_types |= handle_type(subclass)
return contained_special_types
def _make_loader(cls: Type[Any]) -> Type[yaml.Loader]:
import yaml
class DataclassLoader(yaml.Loader):
pass
# Design Q: do we want to support multiple dataclass types with the same name?
# - Why yes: rudimentary support for this is easy.
# - Why no: might cause confusing error messages? Using all possible context cues
# for this is also hard.
#
# => let's just keep things simple, assert uniqueness for now. Easier to add new
# features later than remove them.
contained_types = list(_get_contained_special_types_from_type(cls))
contained_type_names = list(map(lambda cls: cls.__name__, contained_types))
assert len(set(contained_type_names)) == len(contained_type_names), (
"Contained dataclass type names must all be unique, but got"
f" {contained_type_names}"
)
def make_dataclass_constructor(typ: Type[Any]):
return lambda loader, node: typ(**loader.construct_mapping(node))
def make_enum_constructor(typ: Type[enum.Enum]):
return lambda loader, node: typ[loader.construct_python_str(node)]
for typ, name in zip(contained_types, contained_type_names):
if dataclasses.is_dataclass(typ):
DataclassLoader.add_constructor(
tag=DATACLASS_YAML_TAG_PREFIX + name,
constructor=make_dataclass_constructor(typ),
)
elif issubclass(typ, enum.Enum):
DataclassLoader.add_constructor(
tag=ENUM_YAML_TAG_PREFIX + name,
constructor=make_enum_constructor(typ),
)
else:
assert False
DataclassLoader.add_constructor(
tag=MISSING_YAML_TAG_PREFIX,
constructor=lambda *_unused: MISSING, # type: ignore
)
return DataclassLoader
def _make_dumper(instance: Any) -> Type[yaml.Dumper]:
import yaml
class DataclassDumper(yaml.Dumper):
def ignore_aliases(self, data):
return super().ignore_aliases(data) or data is MISSING
contained_types = list(_get_contained_special_types_from_type(type(instance)))
contained_type_names = list(map(lambda cls: cls.__name__, contained_types))
# Note: this is currently a stricter than necessary assert.
assert len(set(contained_type_names)) == len(contained_type_names), (
"Contained dataclass/enum names must all be unique, but got"
f" {contained_type_names}"
)
def make_representer(name: str):
def representer(dumper: DataclassDumper, data: Any) -> yaml.Node:
if dataclasses.is_dataclass(data):
return dumper.represent_mapping(
tag=DATACLASS_YAML_TAG_PREFIX + name,
mapping={
field.name: getattr(data, field.name)
for field in dataclasses.fields(data)
if field.init
},
)
elif isinstance(data, enum.Enum):
return dumper.represent_scalar(
tag=ENUM_YAML_TAG_PREFIX + name, value=data.name
)
assert False
return representer
for typ, name in zip(contained_types, contained_type_names):
DataclassDumper.add_representer(typ, make_representer(name))
DataclassDumper.add_representer(
type(MISSING),
lambda dumper, data: dumper.represent_scalar(
tag=MISSING_YAML_TAG_PREFIX, value=""
),
)
return DataclassDumper
[docs]
def from_yaml(
cls: Type[DataclassType],
stream: Union[str, IO[str], bytes, IO[bytes]],
) -> DataclassType:
"""Re-construct a dataclass instance from a yaml-compatible string, which should be
generated from :func:`tyro.extras.to_yaml()`.
.. warning::
**Deprecated.** Serialization functionality is stable but deprecated.
It may be removed in a future version of :code:`tyro`.
As a secondary feature aimed at enabling the use of :func:`tyro.cli` for general
configuration use cases, we also introduce functions for human-readable dataclass
serialization: :func:`tyro.extras.from_yaml` and :func:`tyro.extras.to_yaml` attempt
to strike a balance between flexibility and robustness — in contrast to naively
dumping or loading dataclass instances (via pickle, PyYAML, etc), explicit type
references enable custom tags that are robust against code reorganization and
refactor, while a PyYAML backend enables serialization of arbitrary Python objects.
Args:
cls: Type to reconstruct.
stream: YAML to read from.
Returns:
Instantiated dataclass.
"""
import yaml
out = yaml.load(stream, Loader=_make_loader(cls))
origin_cls = get_origin(cls)
assert isinstance(out, origin_cls if origin_cls is not None else cls)
return out
[docs]
def to_yaml(instance: Any) -> str:
"""Serialize a dataclass; returns a yaml-compatible string that can be deserialized
via :func:`tyro.extras.from_yaml()`.
.. warning::
**Deprecated.** Serialization functionality is stable but deprecated.
It may be removed in a future version of :code:`tyro`.
As a secondary feature aimed at enabling the use of :func:`tyro.cli` for general
configuration use cases, we also introduce functions for human-readable dataclass
serialization: :func:`tyro.extras.from_yaml` and :func:`tyro.extras.to_yaml` attempt
to strike a balance between flexibility and robustness — in contrast to naively
dumping or loading dataclass instances (via pickle, PyYAML, etc), explicit type
references enable custom tags that are robust against code reorganization and
refactor, while a PyYAML backend enables serialization of arbitrary Python objects.
Args:
instance: Dataclass instance to serialize.
Returns:
YAML string.
"""
import yaml
return "# tyro YAML.\n" + yaml.dump(instance, Dumper=_make_dumper(instance))