55 lines
1.4 KiB
Python
55 lines
1.4 KiB
Python
|
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),
|
||
|
|
)
|