Source code for valkey_dict.type_management

"""Type management module."""

import json
import base64
from collections import OrderedDict, defaultdict
from datetime import datetime, date, time, timedelta

from typing import Callable, Any, Dict, Tuple, Set

from uuid import UUID
from decimal import Decimal


SENTINEL = object()

EncodeFuncType = Callable[[Any], str]
DecodeFuncType = Callable[[str], Any]
EncodeType = Dict[str, EncodeFuncType]
DecodeType = Dict[str, DecodeFuncType]


def _create_default_encode(custom_encode_method: str) -> EncodeFuncType:
    def default_encode(obj: Any) -> str:
        return getattr(obj, custom_encode_method)()  # type: ignore[no-any-return]
    return default_encode


def _create_default_decode(cls: type, custom_decode_method: str) -> DecodeFuncType:
    def default_decode(encoded_str: str) -> Any:
        return getattr(cls, custom_decode_method)(encoded_str)
    return default_decode


def _decode_tuple(val: str) -> Tuple[Any, ...]:
    """
    Deserialize a JSON-formatted string to a tuple.

    This function takes a JSON-formatted string, deserializes it to a list, and
    then converts the list to a tuple.

    Args:
        val (str): A JSON-formatted string representing a list.

    Returns:
        Tuple[Any, ...]: A tuple with the deserialized values from the input string.
    """
    return tuple(json.loads(val))


def _encode_tuple(val: Tuple[Any, ...]) -> str:
    """
    Serialize a tuple to a JSON-formatted string.

    This function takes a tuple, converts it to a list, and then serializes
    the list to a JSON-formatted string.

    Args:
        val (Tuple[Any, ...]): A tuple with values to be serialized.

    Returns:
        str: A JSON-formatted string representing the input tuple.
    """
    return json.dumps(list(val))


def _decode_set(val: str) -> Set[Any]:
    """
    Deserialize a JSON-formatted string to a set.

    This function takes a JSON-formatted string, deserializes it to a list, and
    then converts the list to a set.

    Args:
        val (str): A JSON-formatted string representing a list.

    Returns:
        set[Any]: A set with the deserialized values from the input string.
    """
    return set(json.loads(val))


def _encode_set(val: Set[Any]) -> str:
    """
    Serialize a set to a JSON-formatted string.

    This function takes a set, converts it to a list, and then serializes the
    list to a JSON-formatted string.

    Args:
        val (set[Any]): A set with values to be serialized.

    Returns:
        str: A JSON-formatted string representing the input set.
    """
    return json.dumps(list(val))


decoding_registry: DecodeType = {
    type('').__name__: str,
    type(1).__name__: int,
    type(0.1).__name__: float,
    type(True).__name__: lambda x: x == "True",
    type(None).__name__: lambda x: None,

    "list": json.loads,
    "dict": json.loads,
    "tuple": _decode_tuple,
    type(set()).__name__: _decode_set,

    datetime.__name__: datetime.fromisoformat,
    date.__name__: date.fromisoformat,
    time.__name__: time.fromisoformat,
    timedelta.__name__: lambda x: timedelta(seconds=float(x)),

    Decimal.__name__: Decimal,
    complex.__name__: lambda x: complex(*map(float, x.split(','))),
    bytes.__name__: base64.b64decode,

    UUID.__name__: UUID,
    OrderedDict.__name__: lambda x: OrderedDict(json.loads(x)),
    defaultdict.__name__: lambda x: defaultdict(type(None), json.loads(x)),
    frozenset.__name__: lambda x: frozenset(json.loads(x)),
}


encoding_registry: EncodeType = {
    "list": json.dumps,
    "dict": json.dumps,
    "tuple": _encode_tuple,
    type(set()).__name__: _encode_set,

    datetime.__name__: datetime.isoformat,
    date.__name__: date.isoformat,
    time.__name__: time.isoformat,
    timedelta.__name__: lambda x: str(x.total_seconds()),

    complex.__name__: lambda x: f"{x.real},{x.imag}",
    bytes.__name__: lambda x: base64.b64encode(x).decode('ascii'),
    OrderedDict.__name__: lambda x: json.dumps(list(x.items())),
    defaultdict.__name__: lambda x: json.dumps(dict(x)),
    frozenset.__name__: lambda x: json.dumps(list(x)),
}


[docs] class ValkeyDictJSONEncoder(json.JSONEncoder): """ Extends JSON encoding capabilities by reusing ValkeyDict type conversion. Uses existing decoding_registry to know which types to handle specially and encoding_registry (falls back to str) for converting to JSON-compatible formats. Example: The encoded format looks like:: { "__type__": "TypeName", "value": <encoded value> } Notes: Uses decoding_registry (containing all supported types) to check if type needs special handling. For encoding, defaults to str() if no encoder exists in encoding_registry. """
[docs] def default(self, o: Any) -> Any: """Overwrite default from json encoder. Args: o (Any): Object to be serialized. Raises: TypeError: If the object `o` cannot be serialized. Returns: Any: Serialized value. """ type_name = type(o).__name__ if type_name in decoding_registry: return { "__type__": type_name, "value": encoding_registry.get(type_name, _default_encoder)(o) } try: return json.JSONEncoder.default(self, o) except TypeError as e: raise TypeError(f"Object of type {type_name} is not JSON serializable") from e
[docs] class ValkeyDictJSONDecoder(json.JSONDecoder): """ JSON decoder leveraging ValkeyDict existing type conversion system. Works with ValkeyDictJSONEncoder to reconstruct Python objects from JSON using ValkeyDict decoding_registry. Still needs work but allows for more types than without. """
[docs] def __init__(self, *args: Any, **kwargs: Any) -> None: """ Overwrite the __init__ method from JSON decoder. Args: *args (Any): Positional arguments for initialization. **kwargs (Any): Keyword arguments for initialization. """ def _object_hook(obj: Dict[Any, Any]) -> Any: if "__type__" in obj and "value" in obj: type_name = obj["__type__"] if type_name in decoding_registry: return decoding_registry[type_name](obj["value"]) return obj super().__init__(object_hook=_object_hook, *args, **kwargs)
[docs] def encode_json(obj: Any) -> str: """ Encode a Python object to a JSON string using the existing encoding registry. Args: obj (Any): The Python object to be encoded. Returns: str: The JSON-encoded string representation of the object. """ return json.dumps(obj, cls=ValkeyDictJSONEncoder)
[docs] def decode_json(s: str) -> Any: """ Decode a JSON string to a Python object using the existing decoding registry. Args: s (str): The JSON string to be decoded. Returns: Any: The decoded Python object. """ return json.loads(s, cls=ValkeyDictJSONDecoder)
def _default_decoder(x: str) -> str: """ Pass-through decoder that returns the input string unchanged. Args: x (str): The input string. Returns: str: The same input string. """ return x def _default_encoder(x: Any) -> str: """ Take x and returns the result str of the object. Args: x (Any): The input object Returns: str: output of str of the object """ return str(x) encoding_registry["dict"] = encode_json decoding_registry["dict"] = decode_json encoding_registry["list"] = encode_json decoding_registry["list"] = decode_json