Coverage for qreward / utils / patch.py: 100.00%

1import os (empty)

2 

3import httpx (empty)

4 

5# 可以通过环境变量控制 

6_default_json_lib = os.getenv( (empty)

7 key="JSON_LIB", 

8 default=None, 

9) 

10 

11 

12def patch_httpx(json_lib_name: str = _default_json_lib): # type: ignore (empty)

13 _use_ujson = False (empty)

14 _use_orjson = False (empty)

15 

16 from httpx._content import Any, ByteStream, json_dumps (empty)

17 from httpx._models import typing, jsonlib (empty)

18 

19 if json_lib_name == "ujson": (empty)

20 try: (empty)

21 import ujson (empty)

22 

23 _use_ujson = True (empty)

24 except ImportError: (empty)

25 pass (empty)

26 elif json_lib_name == "orjson": (empty)

27 try: (empty)

28 import orjson (empty)

29 

30 _use_orjson = True (empty)

31 except ImportError: (empty)

32 pass (empty)

33 

34 def encode_json(json: Any) -> tuple[dict[str, str], ByteStream]: (empty)

35 if _use_orjson: (empty)

36 body = orjson.dumps(json) (empty)

37 elif _use_ujson: (empty)

38 body = ujson.dumps( (empty)

39 json, 

40 ensure_ascii=False, 

41 separators=(",", ":"), 

42 allow_nan=False, 

43 ) 

44 else: 

45 body = json_dumps( (empty)

46 json, 

47 ensure_ascii=False, 

48 separators=(",", ":"), 

49 allow_nan=False, 

50 ) 

51 

52 if isinstance(body, str): (empty)

53 body = body.encode("utf-8") (empty)

54 

55 content_length = str(len(body)) (empty)

56 content_type = "application/json" (empty)

57 headers = { (empty)

58 "Content-Length": content_length, 

59 "Content-Type": content_type, 

60 } 

61 return headers, ByteStream(body) (empty)

62 

63 def decode_json(self, **kwargs: typing.Any) -> typing.Any: (empty)

64 if _use_orjson: (empty)

65 return orjson.loads(self.content) (empty)

66 elif _use_ujson: (empty)

67 return ujson.loads(self.content, **kwargs) (empty)

68 return jsonlib.loads(self.content, **kwargs) (empty)

69 

70 # encode 

71 encode_json.__globals__.update(httpx._content.__dict__) (empty)

72 encode_json.__module__ = httpx._content.__name__ (empty)

73 httpx._content.encode_json = encode_json (empty)

74 

75 # decode 

76 httpx._models.Response.json = decode_json (empty)