lodum.yaml

  1# SPDX-FileCopyrightText: 2025-present Michael R. Bernstein <zopemaven@gmail.com>
  2#
  3# SPDX-License-Identifier: Apache-2.0
  4import io
  5from typing import Any, Dict, Iterator, Type, TypeVar
  6
  7try:
  8    from ruamel.yaml import YAML
  9
 10    yaml_available = True
 11except ImportError:
 12    YAML = None  # type: ignore
 13    yaml_available = False
 14
 15from .core import Loader, BaseDumper, BaseLoader
 16from .internal import dump, load, DEFAULT_MAX_SIZE, generate_schema
 17from .exception import DeserializationError
 18
 19T = TypeVar("T")
 20yaml: Any = None
 21if yaml_available:
 22    yaml = YAML(typ="safe")
 23    yaml.sort_base_mapping_type_on_output = False
 24
 25# --- Public API ---
 26
 27
 28def dumps(obj: Any) -> str:
 29    """
 30    Encodes a Python object to a YAML string.
 31
 32    Args:
 33        obj: The object to encode. Must be lodum-enabled or a supported type.
 34
 35    Returns:
 36        A YAML string representation of the object.
 37
 38    Raises:
 39        ImportError: If ruamel.yaml is not installed.
 40    """
 41    if not yaml_available:
 42        raise ImportError(
 43            "ruamel.yaml is required for YAML serialization. Install it with 'pip install lodum[yaml]'."
 44        )
 45
 46    dumper = YamlDumper()
 47    dumped_data = dump(obj, dumper)
 48
 49    with io.StringIO() as string_stream:
 50        yaml.dump(dumped_data, string_stream)
 51        return string_stream.getvalue()
 52
 53
 54def loads(cls: Type[T], yaml_string: str, max_size: int = DEFAULT_MAX_SIZE) -> T:
 55    """
 56    Decodes a YAML string into a Python object of the specified type.
 57
 58    Args:
 59        cls: The class to instantiate.
 60        yaml_string: The YAML data to decode.
 61        max_size: Maximum allowed size of the input string in bytes.
 62
 63    Returns:
 64        An instance of cls populated with the decoded data.
 65
 66    Raises:
 67        DeserializationError: If the input is invalid or exceeds max_size.
 68        ImportError: If ruamel.yaml is not installed.
 69    """
 70    if len(yaml_string) > max_size:
 71        raise DeserializationError(
 72            f"Input size ({len(yaml_string)}) exceeds maximum allowed ({max_size})"
 73        )
 74
 75    if not yaml_available:
 76        raise ImportError(
 77            "ruamel.yaml is required for YAML deserialization. Install it with 'pip install lodum[yaml]'."
 78        )
 79
 80    data = yaml.load(yaml_string)
 81    loader = YamlLoader(data)
 82    return load(cls, loader)
 83
 84
 85def schema(cls: Type[Any]) -> Dict[str, Any]:
 86    """Generates a JSON Schema for a given lodum-enabled class."""
 87    return generate_schema(cls)
 88
 89
 90# --- YAML Dumper Implementation ---
 91
 92
 93class YamlDumper(BaseDumper):
 94    """
 95    Encodes Python objects into a YAML-compatible intermediate representation.
 96    """
 97
 98    def dump_bytes(self, value: bytes) -> Any:
 99        # YAML can handle bytes natively if using certain tags,
100        # but for simplicity and cross-format consistency, we'll use base64 like JSON.
101        import base64
102
103        return base64.b64encode(value).decode("ascii")
104
105
106# --- YAML Loader Implementation ---
107
108
109class YamlLoader(BaseLoader):
110    """
111    Decodes a YAML-compatible intermediate representation into Python objects.
112    """
113
114    def load_list(self) -> Iterator["Loader"]:
115        if not isinstance(self._data, list):
116            raise DeserializationError(
117                f"Expected list, got {type(self._data).__name__}"
118            )
119        return (YamlLoader(item) for item in self._data)
120
121    def load_dict(self) -> Iterator[tuple[str, "Loader"]]:
122        if not isinstance(self._data, dict):
123            raise DeserializationError(
124                f"Expected dict, got {type(self._data).__name__}"
125            )
126        return ((k, YamlLoader(v)) for k, v in self._data.items())
127
128    def load_bytes_value(self, value: Any) -> bytes:
129        if isinstance(value, bytes):
130            return value
131        if not isinstance(value, str):
132            raise DeserializationError(f"Expected str, got {type(value).__name__}")
133        import base64
134
135        try:
136            return base64.b64decode(value)
137        except Exception as e:
138            raise DeserializationError(f"Failed to decode base64: {e}")