Underground-Digital/Workflow-Engine

1import dataclasses2import datetime3from collections import defaultdict, deque4from collections.abc import Callable5from decimal import Decimal6from enum import Enum7from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network8from pathlib import Path, PurePath9from re import Pattern10from types import GeneratorType11from typing import Any, Literal, Optional, Union12from uuid import UUID13 14from pydantic import BaseModel15from pydantic.networks import AnyUrl, NameEmail16from pydantic.types import SecretBytes, SecretStr17from pydantic_core import Url18from pydantic_extra_types.color import Color19 20 21def _model_dump(model: BaseModel, mode: Literal["json", "python"] = "json", **kwargs: Any) -> Any:22    return model.model_dump(mode=mode, **kwargs)23 24 25# Taken from Pydantic v1 as is26def isoformat(o: Union[datetime.date, datetime.time]) -> str:27    return o.isoformat()28 29 30# Taken from Pydantic v1 as is31# TODO: pv2 should this return strings instead?32def decimal_encoder(dec_value: Decimal) -> Union[int, float]:33    """34    Encodes a Decimal as int of there's no exponent, otherwise float35 36    This is useful when we use ConstrainedDecimal to represent Numeric(x,0)37    where a integer (but not int typed) is used. Encoding this as a float38    results in failed round-tripping between encode and parse.39    Our Id type is a prime example of this.40 41    >>> decimal_encoder(Decimal("1.0"))42    1.043 44    >>> decimal_encoder(Decimal("1"))45    146    """47    if dec_value.as_tuple().exponent >= 0:  # type: ignore[operator]48        return int(dec_value)49    else:50        return float(dec_value)51 52 53ENCODERS_BY_TYPE: dict[type[Any], Callable[[Any], Any]] = {54    bytes: lambda o: o.decode(),55    Color: str,56    datetime.date: isoformat,57    datetime.datetime: isoformat,58    datetime.time: isoformat,59    datetime.timedelta: lambda td: td.total_seconds(),60    Decimal: decimal_encoder,61    Enum: lambda o: o.value,62    frozenset: list,63    deque: list,64    GeneratorType: list,65    IPv4Address: str,66    IPv4Interface: str,67    IPv4Network: str,68    IPv6Address: str,69    IPv6Interface: str,70    IPv6Network: str,71    NameEmail: str,72    Path: str,73    Pattern: lambda o: o.pattern,74    SecretBytes: str,75    SecretStr: str,76    set: list,77    UUID: str,78    Url: str,79    AnyUrl: str,80}81 82 83def generate_encoders_by_class_tuples(84    type_encoder_map: dict[Any, Callable[[Any], Any]],85) -> dict[Callable[[Any], Any], tuple[Any, ...]]:86    encoders_by_class_tuples: dict[Callable[[Any], Any], tuple[Any, ...]] = defaultdict(tuple)87    for type_, encoder in type_encoder_map.items():88        encoders_by_class_tuples[encoder] += (type_,)89    return encoders_by_class_tuples90 91 92encoders_by_class_tuples = generate_encoders_by_class_tuples(ENCODERS_BY_TYPE)93 94 95def jsonable_encoder(96    obj: Any,97    by_alias: bool = True,98    exclude_unset: bool = False,99    exclude_defaults: bool = False,100    exclude_none: bool = False,101    custom_encoder: Optional[dict[Any, Callable[[Any], Any]]] = None,102    sqlalchemy_safe: bool = True,103) -> Any:104    custom_encoder = custom_encoder or {}105    if custom_encoder:106        if type(obj) in custom_encoder:107            return custom_encoder[type(obj)](obj)108        else:109            for encoder_type, encoder_instance in custom_encoder.items():110                if isinstance(obj, encoder_type):111                    return encoder_instance(obj)112    if isinstance(obj, BaseModel):113        obj_dict = _model_dump(114            obj,115            mode="json",116            include=None,117            exclude=None,118            by_alias=by_alias,119            exclude_unset=exclude_unset,120            exclude_none=exclude_none,121            exclude_defaults=exclude_defaults,122        )123        if "__root__" in obj_dict:124            obj_dict = obj_dict["__root__"]125        return jsonable_encoder(126            obj_dict,127            exclude_none=exclude_none,128            exclude_defaults=exclude_defaults,129            sqlalchemy_safe=sqlalchemy_safe,130        )131    if dataclasses.is_dataclass(obj):132        obj_dict = dataclasses.asdict(obj)133        return jsonable_encoder(134            obj_dict,135            by_alias=by_alias,136            exclude_unset=exclude_unset,137            exclude_defaults=exclude_defaults,138            exclude_none=exclude_none,139            custom_encoder=custom_encoder,140            sqlalchemy_safe=sqlalchemy_safe,141        )142    if isinstance(obj, Enum):143        return obj.value144    if isinstance(obj, PurePath):145        return str(obj)146    if isinstance(obj, str | int | float | type(None)):147        return obj148    if isinstance(obj, Decimal):149        return format(obj, "f")150    if isinstance(obj, dict):151        encoded_dict = {}152        allowed_keys = set(obj.keys())153        for key, value in obj.items():154            if (155                (not sqlalchemy_safe or (not isinstance(key, str)) or (not key.startswith("_sa")))156                and (value is not None or not exclude_none)157                and key in allowed_keys158            ):159                encoded_key = jsonable_encoder(160                    key,161                    by_alias=by_alias,162                    exclude_unset=exclude_unset,163                    exclude_none=exclude_none,164                    custom_encoder=custom_encoder,165                    sqlalchemy_safe=sqlalchemy_safe,166                )167                encoded_value = jsonable_encoder(168                    value,169                    by_alias=by_alias,170                    exclude_unset=exclude_unset,171                    exclude_none=exclude_none,172                    custom_encoder=custom_encoder,173                    sqlalchemy_safe=sqlalchemy_safe,174                )175                encoded_dict[encoded_key] = encoded_value176        return encoded_dict177    if isinstance(obj, list | set | frozenset | GeneratorType | tuple | deque):178        encoded_list = []179        for item in obj:180            encoded_list.append(181                jsonable_encoder(182                    item,183                    by_alias=by_alias,184                    exclude_unset=exclude_unset,185                    exclude_defaults=exclude_defaults,186                    exclude_none=exclude_none,187                    custom_encoder=custom_encoder,188                    sqlalchemy_safe=sqlalchemy_safe,189                )190            )191        return encoded_list192 193    if type(obj) in ENCODERS_BY_TYPE:194        return ENCODERS_BY_TYPE[type(obj)](obj)195    for encoder, classes_tuple in encoders_by_class_tuples.items():196        if isinstance(obj, classes_tuple):197            return encoder(obj)198 199    try:200        data = dict(obj)201    except Exception as e:202        errors: list[Exception] = []203        errors.append(e)204        try:205            data = vars(obj)206        except Exception as e:207            errors.append(e)208            raise ValueError(errors) from e209    return jsonable_encoder(210        data,211        by_alias=by_alias,212        exclude_unset=exclude_unset,213        exclude_defaults=exclude_defaults,214        exclude_none=exclude_none,215        custom_encoder=custom_encoder,216        sqlalchemy_safe=sqlalchemy_safe,217    )218