turbo_broccoli.user

 1"""User provided encoder/decoder methods."""
 2
 3from typing import Any, Callable, TypeAlias
 4
 5from .context import Context
 6
 7Encoder: TypeAlias = Callable[[Any, Context], dict]
 8Decoder: TypeAlias = Callable[[dict, Context], Any]
 9ClassOrClasses: TypeAlias = (
10    type | str | list[type | str] | tuple[type | str] | set[type | str]
11)
12
13encoders: dict[str, Encoder] = {}
14decoders: dict[str, Decoder] = {}
15
16
17def register_encoder(encoder: Encoder | None, class_or_tuple: ClassOrClasses):
18    """
19    Register a custom encoder for the given type(s). An encoder is a function
20    that takes an object and a `Context` and returns a dict with the following
21    structure:
22
23    ```py
24    {
25        "__type__": "user.<a_type_name>",
26        ...
27    }
28    ```
29
30    `a_type_name` should be a name used when registering the corresponding
31    decoder.
32
33    The dict can be/contain types that are not readily JSON serializable as
34    long as they can be serialized by TurboBroccoli or by other user-provided
35    encoders.
36
37    Args:
38        encoder (Encoder | None): If `None`, the encoder is removed.
39        class_or_tuple (ClassOrClasses): See `ClassOrClasses`
40    """
41    if not isinstance(class_or_tuple, (list, tuple, set)):
42        class_or_tuple = [class_or_tuple]
43    if len(class_or_tuple) == 0:
44        raise ValueError(
45            "At least one class must be provided to register an encoder"
46        )
47    for cls in class_or_tuple:
48        name = cls if isinstance(cls, str) else cls.__name__
49        if encoder is None and name in encoders:
50            del encoders[name]
51        elif encoder is not None:
52            encoders[name] = encoder
53
54
55def register_decoder(decoder: Decoder | None, class_or_tuple: ClassOrClasses):
56    """
57    Register a custom encoder for the given type(s). An decoder is a function
58    that takes a dict and a `Context` and returns an object. The dict will have
59    the form specified in `register_encoder`. It contains data that have
60    already been deserialized by TurboBroccoli or by other user-provided
61    decoders.
62
63    Args:
64        decoder (Decoder | None): If `None`, the decoder is removed.
65        class_or_tuple (ClassOrClasses): See `ClassOrClasses`
66    """
67    if not isinstance(class_or_tuple, (list, tuple, set)):
68        class_or_tuple = [class_or_tuple]
69    if len(class_or_tuple) == 0:
70        raise ValueError(
71            "At least one class must be provided to register a decoder"
72        )
73    for cls in class_or_tuple:
74        name = cls if isinstance(cls, str) else cls.__name__
75        if decoder is None and name in decoders:
76            del decoders[name]
77        elif decoder is not None:
78            decoders[name] = decoder