lodum.bson

  1# SPDX-FileCopyrightText: 2025-present Michael R. Bernstein <zopemaven@gmail.com>
  2#
  3# SPDX-License-Identifier: Apache-2.0
  4try:
  5    import bson
  6except ImportError:
  7    bson = None  # type: ignore
  8from typing import Any, Iterator, Type, TypeVar
  9
 10from .core import Loader, BaseDumper, BaseLoader
 11from .exception import DeserializationError
 12from .internal import dump, load, DEFAULT_MAX_SIZE
 13
 14T = TypeVar("T")
 15
 16# --- Public API ---
 17
 18
 19def dumps(obj: Any) -> bytes:
 20    """
 21    Encodes a Python object to BSON bytes.
 22
 23    Args:
 24        obj: The object to encode. Must be lodum-enabled or a supported type.
 25
 26    Returns:
 27        The BSON-encoded bytes.
 28
 29    Raises:
 30        ImportError: If bson (pymongo) is not installed.
 31    """
 32    if bson is None:
 33        raise ImportError(
 34            "bson (pymongo) is required for BSON serialization. Install it with 'pip install lodum[bson]'."
 35        )
 36    dumper = BsonDumper()
 37    dumped_data = dump(obj, dumper)
 38    # BSON requires a dictionary at the root
 39    if not isinstance(dumped_data, dict):
 40        dumped_data = {"_v": dumped_data}
 41    return bson.encode(dumped_data)
 42
 43
 44def loads(cls: Type[T], bson_bytes: bytes, max_size: int = DEFAULT_MAX_SIZE) -> T:
 45    """
 46    Decodes BSON bytes into a Python object of the specified type.
 47
 48    Args:
 49        cls: The class to instantiate.
 50        bson_bytes: The BSON data to decode.
 51        max_size: Maximum allowed size of the input bytes.
 52
 53    Returns:
 54        An instance of cls populated with the decoded data.
 55
 56    Raises:
 57        DeserializationError: If the input is invalid or exceeds max_size.
 58        ImportError: If bson (pymongo) is not installed.
 59    """
 60    if len(bson_bytes) > max_size:
 61        raise DeserializationError(
 62            f"Input size ({len(bson_bytes)}) exceeds maximum allowed ({max_size})"
 63        )
 64
 65    if bson is None:
 66        raise ImportError(
 67            "bson (pymongo) is required for BSON deserialization. Install it with 'pip install lodum[bson]'."
 68        )
 69    try:
 70        data = bson.decode(bson_bytes)
 71    except Exception as e:
 72        raise DeserializationError(f"Failed to parse BSON: {e}")
 73
 74    # Check if we wrapped a primitive
 75    if "_v" in data and len(data) == 1:
 76        data = data["_v"]
 77
 78    loader = BsonLoader(data)
 79    return load(cls, loader)
 80
 81
 82# --- BSON Dumper Implementation ---
 83
 84
 85class BsonDumper(BaseDumper):
 86    pass
 87
 88
 89# --- BSON Loader Implementation ---
 90
 91
 92class BsonLoader(BaseLoader):
 93    def load_list(self) -> Iterator["Loader"]:
 94        if not isinstance(self._data, list):
 95            raise DeserializationError(
 96                f"Expected list, got {type(self._data).__name__}"
 97            )
 98        return (BsonLoader(item) for item in self._data)
 99
100    def load_dict(self) -> Iterator[tuple[str, "Loader"]]:
101        if not isinstance(self._data, dict):
102            raise DeserializationError(
103                f"Expected dict, got {type(self._data).__name__}"
104            )
105        return ((k, BsonLoader(v)) for k, v in self._data.items())
106
107    def load_bytes_value(self, value: Any) -> bytes:
108        if not isinstance(value, bytes):
109            raise DeserializationError(f"Expected bytes, got {type(value).__name__}")
110        return value