Coverage for langsmith/_internal/_orjson.py: 0%

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 dataclasses 

17 import json 

18 import uuid 

19 from typing import Any, Callable, Optional, Union 

20 

21 DefaultFunc = Optional[Callable[[Any], Any]] 

22 

23 OPT_NON_STR_KEYS = 1 

24 OPT_SERIALIZE_DATACLASS = 2 

25 OPT_SERIALIZE_NUMPY = 4 

26 OPT_SERIALIZE_UUID = 8 

27 

28 class Fragment: # type: ignore 

29 def __init__(self, payloadb: bytes): 

30 self.payloadb = payloadb 

31 

32 from json import JSONDecodeError # type: ignore 

33 

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 orjson 

42 # enable_non_str_keys = bool(option & OPT_NON_STR_KEYS) 

43 if option is None: 

44 option = 0 

45 

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: ignore 

51 def encode(self, o: Any) -> str: 

52 if isinstance(o, Fragment): 

53 return o.payloadb.decode("utf-8") # type: ignore 

54 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() function 

61 return o.tolist() 

62 if ( 

63 enable_serialize_dataclass 

64 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]