codekingpro/portable-devtools

1"""Stubs for orjson operations, compatible with PyPy via a json fallback."""2 3try:4    from orjson import (5        OPT_NON_STR_KEYS,6        OPT_SERIALIZE_DATACLASS,7        OPT_SERIALIZE_NUMPY,8        OPT_SERIALIZE_UUID,9        Fragment,10        JSONDecodeError,11        dumps,12        loads,13    )14 15except ImportError:16    import dataclasses17    import json18    import uuid19    from typing import Any, Callable, Optional, Union20 21    DefaultFunc = Optional[Callable[[Any], Any]]22 23    OPT_NON_STR_KEYS = 124    OPT_SERIALIZE_DATACLASS = 225    OPT_SERIALIZE_NUMPY = 426    OPT_SERIALIZE_UUID = 827 28    class Fragment:  # type: ignore29        def __init__(self, payloadb: bytes):30            self.payloadb = payloadb31 32    from json import JSONDecodeError  # type: ignore33 34    def dumps(35        obj: Any,36        /,37        default: DefaultFunc = None,38        option: Optional[int] = None,39    ) -> bytes:40        # for now, don't do anything for this case because `json.dumps`41        # automatically encodes non-str keys as str by default, unlike orjson42        # enable_non_str_keys = bool(option & OPT_NON_STR_KEYS)43        if option is None:44            option = 045 46        enable_serialize_numpy = bool(option & OPT_SERIALIZE_NUMPY)47        enable_serialize_dataclass = bool(option & OPT_SERIALIZE_DATACLASS)48        enable_serialize_uuid = bool(option & OPT_SERIALIZE_UUID)49 50        class CustomEncoder(json.JSONEncoder):  # type: ignore51            def encode(self, o: Any) -> str:52                if isinstance(o, Fragment):53                    return o.payloadb.decode("utf-8")  # type: ignore54                return super().encode(o)55 56            def default(self, o: Any) -> Any:57                if enable_serialize_uuid and isinstance(o, uuid.UUID):58                    return str(o)59                if enable_serialize_numpy and hasattr(o, "tolist"):60                    # even objects like np.uint16(15) have a .tolist() function61                    return o.tolist()62                if (63                    enable_serialize_dataclass64                    and dataclasses.is_dataclass(o)65                    and not isinstance(o, type)66                ):67                    return dataclasses.asdict(o)68                if default is not None:69                    return default(o)70 71                return super().default(o)72 73        return json.dumps(obj, cls=CustomEncoder).encode("utf-8")74 75    def loads(payload: Union[bytes, bytearray, memoryview, str], /) -> Any:76        return json.loads(payload)77 78 79__all__ = [80    "loads",81    "dumps",82    "Fragment",83    "JSONDecodeError",84    "OPT_SERIALIZE_NUMPY",85    "OPT_SERIALIZE_DATACLASS",86    "OPT_SERIALIZE_UUID",87    "OPT_NON_STR_KEYS",88]89