this repo has no description
dekker.one
1import yaml
2
3try:
4 from yaml import CLoader as Loader, CDumper as Dumper
5except ImportError:
6 from yaml import Loader, Dumper # type: ignore
7
8import datetime
9
10
11class Undefined:
12 """
13 Represents missing values in YAML (as opposed to null which is None in Python)
14 """
15
16 pass
17
18
19def mapping(tag):
20 """
21 A decorator which allows for serializing/deserializing a class as a YAML mapping (dictionary).
22 The class must be able to be constructed with `**kwargs` to set its data.
23
24 :param tag: The tag to use for the type (e.g. `!MyTag`)
25 :type tag: str
26 """
27
28 def decorator(obj_class):
29 def construct(loader, node):
30 try:
31 attrs = loader.construct_mapping(node)
32 instance = obj_class(**attrs)
33 return instance
34 except yaml.constructor.ConstructorError:
35 return obj_class()
36
37 def represent(dumper, data):
38 attrs = {
39 key: value
40 for key, value in data.__dict__.items()
41 if value is not Undefined
42 }
43 return dumper.represent_mapping(tag, attrs)
44
45 yaml.add_constructor(tag, construct, Loader=Loader)
46 yaml.add_representer(obj_class, represent, Dumper=Dumper)
47 return obj_class
48
49 return decorator
50
51
52def sequence(tag):
53 """
54 A decorator which allows for serializing/deserializing a class as a YAML sequence (list).
55 The class must be able to be constructed with `*args` to set it the items, and also be iterable.
56
57 :param tag: The tag to use for the type (e.g. `!MyTag`)
58 :type tag: str
59 """
60
61 def decorator(obj_class):
62 def construct(loader, node):
63 try:
64 args = loader.construct_sequence(node)
65 return obj_class(*args)
66 except yaml.constructor.ConstructorError:
67 return obj_class()
68
69 def represent(dumper, data):
70 return dumper.represent_sequence(tag, iter(data))
71
72 yaml.add_constructor(tag, construct, Loader=Loader)
73 yaml.add_representer(obj_class, represent, Dumper=Dumper)
74 return obj_class
75
76 return decorator
77
78
79def scalar(tag):
80 """
81 A decorator which allows for serializing/deserializing a class as a YAML scalar (string).
82 The class must be able to be constructed with a string argument, and implement `get_value()`
83 to return the string to be serialized.
84
85 :param tag: The tag to use for the type (e.g. `!MyTag`)
86 :type tag: str
87 """
88
89 def decorator(obj_class):
90 def construct(loader, node):
91 value = loader.construct_scalar(node)
92 return obj_class(value)
93
94 def represent(dumper, data):
95 return dumper.represent_scalar(tag, str(data.get_value()))
96
97 yaml.add_constructor(tag, construct, Loader=Loader)
98 yaml.add_representer(obj_class, represent, Dumper=Dumper)
99 return obj_class
100
101 return decorator
102
103
104def load(stream):
105 """
106 Helper function which loads YAML
107 """
108 return yaml.load(stream, Loader=Loader)
109
110
111def load_all(stream):
112 """
113 Helper function which loads YAML as a list of documents
114 """
115 return yaml.load_all(stream, Loader=Loader)
116
117
118def dump(data):
119 """
120 Helper function which serializes objects as YAML
121 """
122 return yaml.dump(data, Dumper=Dumper)
123
124
125def dump_all(data):
126 """
127 Helper function which serializes a list of objects as YAML
128 """
129 return yaml.dump_all(data, Dumper=Dumper)
130
131
132def range_representer(dumper, data):
133 """
134 A YAML `!Range l..u` tag
135 """
136 scalar = u"{}..{}".format(data.start, data.stop - 1)
137 return dumper.represent_scalar(u"!Range", scalar)
138
139
140yaml.add_representer(range, range_representer, Dumper=Dumper)
141
142
143def range_constructor(loader, node):
144 """
145 A YAML `!Range l..u` tag
146 """
147 value = loader.construct_scalar(node)
148 a, b = map(int, value.split(".."))
149 return range(a, b + 1)
150
151
152yaml.add_constructor(u"!Range", range_constructor, Loader=Loader)
153
154
155def dt_representer(dumper, data):
156 """
157 A YAML `!Duration` tag
158 """
159 scalar = u"{}ms".format(data.total_seconds() * 1000)
160 return dumper.represent_scalar(u"!Duration", scalar)
161
162
163yaml.add_representer(datetime.timedelta, dt_representer, Dumper=Dumper)
164
165
166def dt_constructor(loader, node):
167 """
168 A YAML `!Duration` tag
169 """
170 value = loader.construct_scalar(node)
171 if value.endswith("us"):
172 return datetime.timedelta(microseconds=float(value[:-2]))
173 if value.endswith("ms"):
174 return datetime.timedelta(milliseconds=float(value[:-2]))
175 if value.endswith("s"):
176 return datetime.timedelta(seconds=float(value[:-1]))
177
178
179yaml.add_constructor(u"!Duration", dt_constructor, Loader=Loader)
180
181
182def list_representer(dumper, data):
183 """
184 Changes flow style of lists to be on a single line if only primitives are contained
185 """
186 flow_style = all(isinstance(i, (int, float, str, bool)) for i in data)
187 return dumper.represent_sequence(
188 "tag:yaml.org,2002:seq", data, flow_style=flow_style
189 )
190
191
192yaml.add_representer(list, list_representer)