Coverage for o2/util/custom_dumper.py: 87%

1from typing import Union 

2 

3from dataclass_wizard import DumpMixin, LoadMixin 

4 

5from o2.models.timetable.time_period import TimePeriod 

6 

7 

8class CustomDumper(DumpMixin): 

9 """Handles custom serialization of complex data types to dictionaries. 

10 

11 Provides specialized methods for converting complex objects like TimePeriod 

12 into dictionary representations suitable for serialization. 

13 """ 

14 

15 def __init_subclass__(cls) -> None: 

16 """Initialize a subclass with custom dump hooks. 

17 

18 Registers custom dump hooks for complex types that require special handling. 

19 """ 

20 super().__init_subclass__() 

21 # register dump hooks for custom types - used when `to_dict()` is called 

22 cls.register_dump_hook(TimePeriod, cls.dump_with_time_period) 

23 

24 @staticmethod 

25 def dump_with_time_period(o: TimePeriod, *_: object) -> dict: 

26 """Convert a TimePeriod object to a dictionary. 

27 

28 Handles the serialization of TimePeriod objects using their model_dump method. 

29 """ 

30 return o.model_dump(by_alias=True) 

31 

32 

33class CustomLoader(LoadMixin): 

34 """Handles custom deserialization from dictionaries to complex data types. 

35 

36 Provides specialized methods for converting dictionary representations 

37 back into complex objects like TimePeriod. 

38 """ 

39 

40 def __init_subclass__(cls) -> None: 

41 """Initialize a subclass with custom load hooks. 

42 

43 Registers custom load hooks for complex types that require special handling. 

44 """ 

45 super().__init_subclass__() 

46 # register load hooks for custom types - used when `from_dict()` is called 

47 cls.register_load_hook(TimePeriod, cls.load_to_time_period) 

48 

49 @staticmethod 

50 def load_to_time_period(d: Union[str, TimePeriod, dict], base_type: type[TimePeriod]) -> TimePeriod: 

51 """Convert a dictionary, string, or TimePeriod to a TimePeriod object. 

52 

53 Handles deserialization of TimePeriod objects from various input formats. 

54 """ 

55 if isinstance(d, base_type): 

56 return d 

57 if isinstance(d, str): 

58 return TimePeriod.model_validate_json(d) 

59 

60 if isinstance(d, dict): 

61 return TimePeriod.model_validate(d) 

62 

63 raise ValueError(f"Can't load {d} to TimePeriod")