from types import MappingProxyType from typing import Any, Dict, Mapping, Type, TypeVar, Union import yaml from typing_extensions import Protocol from mashumaro.serializer.base import DataClassDictMixin DEFAULT_DICT_PARAMS = { "use_bytes": False, "use_enum": False, "use_datetime": False, } EncodedData = Union[str, bytes] T = TypeVar("T", bound="DataClassYAMLMixin") class Encoder(Protocol): # pragma no cover def __call__(self, o, **kwargs) -> EncodedData: ... class Decoder(Protocol): # pragma no cover def __call__(self, packed: EncodedData, **kwargs) -> Dict[Any, Any]: ... class DataClassYAMLMixin(DataClassDictMixin): __slots__ = () def to_yaml( self: T, encoder: Encoder = yaml.dump, # type: ignore dict_params: Mapping = MappingProxyType({}), **encoder_kwargs, ) -> EncodedData: return encoder( self.to_dict(**dict(DEFAULT_DICT_PARAMS, **dict_params)), **encoder_kwargs, ) @classmethod def from_yaml( cls: Type[T], data: EncodedData, decoder: Decoder = yaml.safe_load, # type: ignore dict_params: Mapping = MappingProxyType({}), **decoder_kwargs, ) -> T: return cls.from_dict( decoder(data, **decoder_kwargs), **dict(DEFAULT_DICT_PARAMS, **dict_params), )