Coverage for src / gitq / yaml.py: 96%

1from typing import Self 

2from dataclasses import fields, is_dataclass, Field 

3 

4import yaml 

5 

6 

7class YAMLObjectMetaclass(yaml.YAMLObjectMetaclass): 

8 

9 def __init__(cls, name, bases, kwds): 

10 cls.yaml_tag = "!" + name 

11 kwds["yaml_tag"] = cls.yaml_tag 

12 super().__init__(name, bases, kwds) 

13 

14 

15def yaml_excluded_fields(cls) -> set[str]: 

16 "Return the set of attribute names marked yaml_exclude via dataclasses.field()." 

17 excluded = set() 

18 for klass in cls.__mro__: 

19 for attr_name, attr_val in vars(klass).items(): 

20 if isinstance(attr_val, Field) and attr_val.metadata.get("yaml_exclude"): 

21 excluded.add(attr_name) 

22 return excluded 

23 

24 

25def represent_value(value, dumper: yaml.Dumper): 

26 if isinstance(value, str) and "\n" in value: 

27 return dumper.represent_scalar("tag:yaml.org,2002:str", value, style="|") 

28 return dumper.represent_data(value) 

29 

30 

31class YAMLObject(yaml.YAMLObject, metaclass=YAMLObjectMetaclass): 

32 

33 # Override to_yaml to customize the yaml representation. 

34 # * Order of fields is as declared in the dataclass 

35 # * Fields marked yaml_exclude are skipped. 

36 # * False values are skipped. 

37 # * Multiline strings are represented with pipe-style yaml strings. 

38 @classmethod 

39 def to_yaml(cls, dumper: yaml.Dumper, data: Self): 

40 excluded = yaml_excluded_fields(cls) 

41 assert is_dataclass(cls), cls.__name__ 

42 

43 def i(): 

44 for f in fields(cls): 

45 if f.name in excluded: 45 ↛ 46line 45 didn't jump to line 46 because the condition on line 45 was never true

46 continue 

47 value = getattr(data, f.name) 

48 if value is None: 

49 continue 

50 yield (dumper.represent_data(f.name), represent_value(value, dumper)) 

51 

52 return yaml.MappingNode(cls.yaml_tag, list(i())) 

53 

54 

55class BaseLoader(yaml.SafeLoader): 

56 

57 # By default, PyYAML uses __new__() and .__dict__.update() to construct 

58 # objects. Use the constructor provided by dataclasses instead, so that 

59 # defaults are respected and unknown fields raise exceptions. 

60 def construct_yaml_object(self, node, cls): 

61 state = self.construct_mapping(node, deep=True) 

62 return cls(**state) # type: ignore