codekingpro/portable-devtools

1from __future__ import annotations2 3import base644import collections5import datetime6import decimal7import ipaddress8import json9import logging10import pathlib11import re12import uuid13from typing import Any14 15from langsmith._internal import _orjson16 17try:18    from zoneinfo import ZoneInfo  # type: ignore[import-not-found]19except ImportError:20 21    class ZoneInfo:  # type: ignore[no-redef]22        """Introduced in python 3.9."""23 24 25logger = logging.getLogger(__name__)26 27 28def _simple_default(obj):29    try:30        # Only need to handle types that orjson doesn't serialize by default31        # https://github.com/ijl/orjson#serialize32        if isinstance(obj, datetime.datetime):33            return obj.isoformat()34        elif isinstance(obj, uuid.UUID):35            return str(obj)36        elif isinstance(obj, BaseException):37            return {"error": type(obj).__name__, "message": str(obj)}38        elif isinstance(obj, (set, frozenset, collections.deque)):39            return list(obj)40        elif isinstance(obj, (datetime.timezone, ZoneInfo)):41            return obj.tzname(None)42        elif isinstance(obj, datetime.timedelta):43            return obj.total_seconds()44        elif isinstance(obj, decimal.Decimal):45            if obj.as_tuple().exponent >= 0:46                return int(obj)47            else:48                return float(obj)49        elif isinstance(50            obj,51            (52                ipaddress.IPv4Address,53                ipaddress.IPv4Interface,54                ipaddress.IPv4Network,55                ipaddress.IPv6Address,56                ipaddress.IPv6Interface,57                ipaddress.IPv6Network,58                pathlib.Path,59            ),60        ):61            return str(obj)62        elif isinstance(obj, re.Pattern):63            return obj.pattern64        elif isinstance(obj, (bytes, bytearray)):65            return base64.b64encode(obj).decode()66        return str(obj)67    except BaseException as e:68        logger.debug(f"Failed to serialize {type(obj)} to JSON: {e}")69    return str(obj)70 71 72_serialization_methods: list[tuple[str, dict[str, Any]]] = [73    (74        "model_dump",75        {"exclude_none": True, "mode": "json"},76    ),  # Pydantic V2 with non-serializable fields77    ("model_dump", {"exclude_none": True}),  # Pydantic V2 without json mode78    ("dict", {}),  # Pydantic V1 with non-serializable field79    ("to_dict", {}),  # dataclasses-json80]81 82 83# IMPORTANT: This function is used from Rust code in `langsmith-pyo3` serialization,84#            in order to handle serializing these tricky Python types *from Rust*.85#            Do not cause this function to become inaccessible (e.g. by deleting86#            or renaming it) without also fixing the corresponding Rust code found in:87#               rust/crates/langsmith-pyo3/src/serialization/mod.rs88def _serialize_json(obj: Any) -> Any:89    try:90        if isinstance(obj, (set, tuple)):91            if hasattr(obj, "_asdict") and callable(obj._asdict):92                # NamedTuple93                return obj._asdict()94            return list(obj)95 96        for attr, kwargs in _serialization_methods:97            if (98                hasattr(obj, attr)99                and callable(getattr(obj, attr))100                and not isinstance(obj, type)101            ):102                try:103                    method = getattr(obj, attr)104                    response = method(**kwargs)105                    if not isinstance(response, dict):106                        return str(response)107                    return response108                except Exception as e:109                    logger.debug(110                        f"Failed to use {attr} to serialize {type(obj)} to"111                        f" JSON: {repr(e)}"112                    )113                    pass114        return _simple_default(obj)115    except BaseException as e:116        logger.debug(f"Failed to serialize {type(obj)} to JSON: {e}")117        return str(obj)118 119 120def _elide_surrogates(s: bytes) -> bytes:121    pattern = re.compile(rb"\\ud[89a-f][0-9a-f]{2}", re.IGNORECASE)122    result = pattern.sub(b"", s)123    return result124 125 126def dumps_json(obj: Any) -> bytes:127    """Serialize an object to a JSON formatted string.128 129    Parameters130    ----------131    obj : Any132        The object to serialize.133    default : Callable[[Any], Any] or None, default=None134        The default function to use for serialization.135 136    Returns:137    -------138    str139        The JSON formatted string.140    """141    try:142        return _orjson.dumps(143            obj,144            default=_serialize_json,145            option=_orjson.OPT_SERIALIZE_NUMPY146            | _orjson.OPT_SERIALIZE_DATACLASS147            | _orjson.OPT_SERIALIZE_UUID148            | _orjson.OPT_NON_STR_KEYS,149        )150    except TypeError as e:151        # Usually caused by UTF surrogate characters152        logger.debug(f"Orjson serialization failed: {repr(e)}. Falling back to json.")153        result = json.dumps(154            obj,155            default=_serialize_json,156            ensure_ascii=True,157        ).encode("utf-8")158        try:159            result = _orjson.dumps(160                _orjson.loads(result.decode("utf-8", errors="surrogateescape"))161            )162        except _orjson.JSONDecodeError:163            result = _elide_surrogates(result)164        return result165