Custom JSON Encoder function

json_encoding.py

"""
Custom JSON Encoding

Thanks Pydantic! https://github.com/samuelcolvin/pydantic/blob/master/pydantic/json.py
"""

from collections import deque
from dataclasses import asdict, is_dataclass
import datetime
from decimal import Decimal
from enum import Enum
from ipaddress import (IPv4Address, IPv4Interface,
                       IPv4Network, IPv6Address,
                       IPv6Interface, IPv6Network)
from pathlib import Path
from re import Pattern
from types import GeneratorType
from typing import Any, Callable, Dict, Type, Union
from uuid import UUID


def isoformat(o: Union[datetime.date, datetime.time]) -> str:
    """
    Converts a Datetime to ISO Format
    """
    return o.isoformat()


def decimal_encoder(dec_value: Decimal) -> Union[int, float]:
    """
    Encodes a Decimal as int of there's no exponent, otherwise float
    """
    if dec_value.as_tuple().exponent >= 0:
        return int(dec_value)
    else:
        return float(dec_value)


# Dictionary of Python type keys and callable values that coerce them to be JSON encoded
ENCODERS_BY_TYPE: Dict[Type[Any], Callable[[Any], Any]] = {
    bytes: lambda b: b.decode(),
    datetime.date: isoformat,
    datetime.datetime: isoformat,
    datetime.time: isoformat,
    datetime.timedelta: str,
    Decimal: decimal_encoder,
    Enum: lambda e: e.value,
    frozenset: list,
    deque: list,
    GeneratorType: list,
    IPv4Address: str,
    IPv4Interface: str,
    IPv4Network: str,
    IPv6Address: str,
    IPv6Interface: str,
    IPv6Network: str,
    Path: str,
    Pattern: lambda p: p.pattern,
    set: list,
    UUID: str,
}

# These Can't Be Imported - So We Call them By Name
ENCODERS_BY_CLASS_NAME: Dict[str, Callable[[Any], Any]] = {
    # NumPy
    "int64": int,
    "int32": int,
    "int16": int,
    "int8": int,
    "float128": float,
    "float64": float,
    "float32": float,
    "float16": float,
    # Pandas
    "DataFrame": lambda df: df.to_dict(orient="records"),
    "Series": list,
    # Pydantic
    "BaseModel": lambda bm: bm.dict(),
}


def custom_json_encoder(obj: Any) -> Any:
    """
    Custom JSON Encoder

    This encoder function checks the class type and its superclasses for a matching encoder

    Parameters
    ----------
    obj : Any
        Object to JSON encode

    Returns
    -------
    Any

    Examples
    ---------
    >>> dict_to_encode = {"datetime": datetime.datetime.now()}
    >>> json.dumps(dict_to_encode, default=custom_json_encoder)
    '{"datetime": "2022-04-13T09:19:29.196264"}'
    """
    if is_dataclass(obj):
        return asdict(obj)
    for base in obj.__class__.__mro__[:-1]:
        try:
            encoder = ENCODERS_BY_TYPE[base]
        except KeyError:
            try:
                class_type = base.__name__
                encoder = ENCODERS_BY_CLASS_NAME[class_type]
            except KeyError:
                continue
        return encoder(obj)
    else:  # We have exited the for loop without finding a suitable encoder
        raise TypeError(f"Object of type '{obj.__class__.__name__}' is not JSON serializable")