Coverage for src/mlopus/utils/yaml.py: 80%

1import functools 

2from dataclasses import is_dataclass, asdict 

3from datetime import datetime 

4from pathlib import Path 

5from typing import Any 

6 

7import yaml 

8 

9from mlopus.utils import typing_utils, import_utils, pydantic, time_utils 

10 

11loads = functools.partial(yaml.load, Loader=yaml.Loader) 

12 

13 

14class Dumper(yaml.Dumper): 

15 """Dumper that handles multiline strings, sets, types, dataclasses, datetime, Path and pydantic objects.""" 

16 

17 def represent_data(self, data): 

18 """Dumper that handles multiline strings, types, dataclasses, datetime and pydantic objects.""" 

19 if isinstance(data, str) and "\n" in data: 

20 data = "\n".join([line.rstrip() for line in data.split("\n")]) 

21 return self.represent_scalar("tag:yaml.org,2002:str", data, style="|") 

22 

23 elif isinstance(data, BaseException): 

24 data = {"type": import_utils.fq_name(type(data)), "message": str(data)} 

25 

26 elif type_ := typing_utils.as_type(data): 

27 data = import_utils.fq_name(type_) 

28 

29 elif p_obj := pydantic.as_model_obj(data): 

30 data = p_obj.dict() 

31 

32 elif is_dataclass(type(data)): 

33 data = asdict(data) 

34 

35 elif isinstance(data, datetime): 

36 data = time_utils.safe_repr(data) 

37 

38 elif isinstance(data, set): 

39 data = list(data) 

40 

41 elif isinstance(data, Path): 

42 data = str(data) 

43 

44 return super().represent_data(data) 

45 

46 

47def dumps(data: Any) -> str: 

48 """Dumper that handles multiline strings, types, dataclasses, datetime and pydantic objects.""" 

49 return yaml.dump(data, sort_keys=False, Dumper=Dumper)