Coverage for mindsdb / utilities / json_encoder.py: 90%

1from datetime import datetime, date, timedelta 

2from decimal import Decimal 

3import pandas as pd 

4import numpy as np 

5import orjson 

6from flask.json.provider import DefaultJSONProvider 

7 

8 

9class CustomJSONEncoder: 

10 def default(self, obj): 

11 if isinstance(obj, timedelta): 

12 return str(obj) 

13 if isinstance(obj, datetime): 

14 return obj.strftime("%Y-%m-%d %H:%M:%S.%f") 

15 if isinstance(obj, date): 

16 return obj.strftime("%Y-%m-%d") 

17 if isinstance(obj, Decimal): 

18 return float(obj) 

19 if isinstance(obj, np.bool_): 19 ↛ 20line 19 didn't jump to line 20 because the condition on line 19 was never true

20 return bool(obj) 

21 if isinstance(obj, np.ndarray): 

22 return obj.tolist() 

23 if pd.isnull(obj): 23 ↛ 26line 23 didn't jump to line 26 because the condition on line 23 was always true

24 return None 

25 

26 return str(obj) 

27 

28 

29class ORJSONProvider(DefaultJSONProvider): 

30 """ 

31 Use orjson to serialize data instead of flask json provider. 

32 """ 

33 

34 def dumps(self, obj, **kwargs): 

35 return orjson.dumps( 

36 obj, 

37 option=( 

38 orjson.OPT_SERIALIZE_NUMPY 

39 | orjson.OPT_NON_STR_KEYS 

40 # keep this for using CustomJSON encoder 

41 | orjson.OPT_PASSTHROUGH_DATETIME 

42 ), 

43 default=CustomJSONEncoder().default, 

44 ).decode("utf-8") 

45 

46 def loads(self, s, **kwargs): 

47 return orjson.loads(s)