Coverage for src/mlopus/utils/json_utils.py: 76%

1import json 

2from dataclasses import is_dataclass, asdict 

3from datetime import datetime, date 

4from enum import Enum 

5from pathlib import Path 

6from typing import Any, Callable, Dict 

7 

8import pydantic 

9 

10from mlopus.utils import time_utils 

11 

12 

13def fallback(obj: Any) -> Dict[str, Any] | list | str | None: 

14 """json.dumps fallback for types, dataclasses, pydantic, date and datetime.""" 

15 match obj: 

16 case Path(): 

17 return str(obj) 

18 case set(): 

19 return sorted(obj) 

20 case type(): 

21 return f"{obj.__module__}.{obj.__class__}" 

22 case pydantic.BaseModel(): 

23 return obj.dict() 

24 case _ if is_dataclass(obj): 

25 return asdict(obj) 

26 case datetime(): 

27 # Get UTC timestamp in full precision ISO format with timezone info: 2024-04-01T09:00:00.000000+00:00 

28 # This preserves datetime sorting when applying alphanumerical sorting to string representation 

29 return time_utils.safe_repr(obj) 

30 case date(): 

31 return obj.isoformat() 

32 case Enum(): 

33 return obj.value 

34 

35 

36def dumps(obj: Any, fallback_: Callable[[Any], Dict[str, Any] | str | None] = fallback, **kwargs) -> str: 

37 """JSON encode with `fallback` function to handle encoding of special types.""" 

38 return json.dumps(obj, default=fallback_, **kwargs) 

39 

40 

41class Decoder(json.JSONDecoder): 

42 """Custom JSON decoder. Handles datetime strings in safe representation (see `time_utils.safe_repr`).""" 

43 

44 def decode(self, s, *args, **kwargs): 

45 """Decode JSON value.""" 

46 return time_utils.maybe_parse_safe_repr(super().decode(s, *args, **kwargs)) 

47 

48 

49def loads(data: str, ignore_errors: bool = False, cls=Decoder) -> Any: 

50 """JSON decode.""" 

51 try: 

52 return json.loads(data, cls=cls) 

53 except json.JSONDecodeError as error: 

54 if ignore_errors: 

55 return str(data) 

56 raise error