| 1 | # SPDX-License-Identifier: MIT |
| 2 | |
| 3 | __all__ = ['BaseRepresenter', 'SafeRepresenter', 'Representer', |
| 4 | 'RepresenterError'] |
| 5 | |
| 6 | from .error import * |
| 7 | from .nodes import * |
| 8 | |
| 9 | import datetime, sys, copyreg, types, base64 |
| 10 | |
| 11 | class RepresenterError(YAMLError): |
| 12 | pass |
| 13 | |
| 14 | class BaseRepresenter: |
| 15 | |
| 16 | yaml_representers = {} |
| 17 | yaml_multi_representers = {} |
| 18 | |
| 19 | def __init__(self, default_style=None, default_flow_style=None): |
| 20 | self.default_style = default_style |
| 21 | self.default_flow_style = default_flow_style |
| 22 | self.represented_objects = {} |
| 23 | self.object_keeper = [] |
| 24 | self.alias_key = None |
| 25 | |
| 26 | def represent(self, data): |
| 27 | node = self.represent_data(data) |
| 28 | self.serialize(node) |
| 29 | self.represented_objects = {} |
| 30 | self.object_keeper = [] |
| 31 | self.alias_key = None |
| 32 | |
| 33 | def represent_data(self, data): |
| 34 | if self.ignore_aliases(data): |
| 35 | self.alias_key = None |
| 36 | else: |
| 37 | self.alias_key = id(data) |
| 38 | if self.alias_key is not None: |
| 39 | if self.alias_key in self.represented_objects: |
| 40 | node = self.represented_objects[self.alias_key] |
| 41 | #if node is None: |
| 42 | # raise RepresenterError("recursive objects are not allowed: %r" % data) |
| 43 | return node |
| 44 | #self.represented_objects[alias_key] = None |
| 45 | self.object_keeper.append(data) |
| 46 | data_types = type(data).__mro__ |
| 47 | if data_types[0] in self.yaml_representers: |
| 48 | node = self.yaml_representers[data_types[0]](self, data) |
| 49 | else: |
| 50 | for data_type in data_types: |
| 51 | if data_type in self.yaml_multi_representers: |
| 52 | node = self.yaml_multi_representers[data_type](self, data) |
| 53 | break |
| 54 | else: |
| 55 | if None in self.yaml_multi_representers: |
| 56 | node = self.yaml_multi_representers[None](self, data) |
| 57 | elif None in self.yaml_representers: |
| 58 | node = self.yaml_representers[None](self, data) |
| 59 | else: |
| 60 | node = ScalarNode(None, str(data)) |
| 61 | #if alias_key is not None: |
| 62 | # self.represented_objects[alias_key] = node |
| 63 | return node |
| 64 | |
| 65 | @classmethod |
| 66 | def add_representer(cls, data_type, representer): |
| 67 | if not 'yaml_representers' in cls.__dict__: |
| 68 | cls.yaml_representers = cls.yaml_representers.copy() |
| 69 | cls.yaml_representers[data_type] = representer |
| 70 | |
| 71 | @classmethod |
| 72 | def add_multi_representer(cls, data_type, representer): |
| 73 | if not 'yaml_multi_representers' in cls.__dict__: |
| 74 | cls.yaml_multi_representers = cls.yaml_multi_representers.copy() |
| 75 | cls.yaml_multi_representers[data_type] = representer |
| 76 | |
| 77 | def represent_scalar(self, tag, value, style=None): |
| 78 | if style is None: |
| 79 | style = self.default_style |
| 80 | node = ScalarNode(tag, value, style=style) |
| 81 | if self.alias_key is not None: |
| 82 | self.represented_objects[self.alias_key] = node |
| 83 | return node |
| 84 | |
| 85 | def represent_sequence(self, tag, sequence, flow_style=None): |
| 86 | value = [] |
| 87 | node = SequenceNode(tag, value, flow_style=flow_style) |
| 88 | if self.alias_key is not None: |
| 89 | self.represented_objects[self.alias_key] = node |
| 90 | best_style = True |
| 91 | for item in sequence: |
| 92 | node_item = self.represent_data(item) |
| 93 | if not (isinstance(node_item, ScalarNode) and not node_item.style): |
| 94 | best_style = False |
| 95 | value.append(node_item) |
| 96 | if flow_style is None: |
| 97 | if self.default_flow_style is not None: |
| 98 | node.flow_style = self.default_flow_style |
| 99 | else: |
| 100 | node.flow_style = best_style |
| 101 | return node |
| 102 | |
| 103 | def represent_mapping(self, tag, mapping, flow_style=None): |
| 104 | value = [] |
| 105 | node = MappingNode(tag, value, flow_style=flow_style) |
| 106 | if self.alias_key is not None: |
| 107 | self.represented_objects[self.alias_key] = node |
| 108 | best_style = True |
| 109 | if hasattr(mapping, 'items'): |
| 110 | mapping = list(mapping.items()) |
| 111 | try: |
| 112 | mapping = sorted(mapping) |
| 113 | except TypeError: |
| 114 | pass |
| 115 | for item_key, item_value in mapping: |
| 116 | node_key = self.represent_data(item_key) |
| 117 | node_value = self.represent_data(item_value) |
| 118 | if not (isinstance(node_key, ScalarNode) and not node_key.style): |
| 119 | best_style = False |
| 120 | if not (isinstance(node_value, ScalarNode) and not node_value.style): |
| 121 | best_style = False |
| 122 | value.append((node_key, node_value)) |
| 123 | if flow_style is None: |
| 124 | if self.default_flow_style is not None: |
| 125 | node.flow_style = self.default_flow_style |
| 126 | else: |
| 127 | node.flow_style = best_style |
| 128 | return node |
| 129 | |
| 130 | def ignore_aliases(self, data): |
| 131 | return False |
| 132 | |
| 133 | class SafeRepresenter(BaseRepresenter): |
| 134 | |
| 135 | def ignore_aliases(self, data): |
| 136 | if data in [None, ()]: |
| 137 | return True |
| 138 | if isinstance(data, (str, bytes, bool, int, float)): |
| 139 | return True |
| 140 | |
| 141 | def represent_none(self, data): |
| 142 | return self.represent_scalar('tag:yaml.org,2002:null', 'null') |
| 143 | |
| 144 | def represent_str(self, data): |
| 145 | return self.represent_scalar('tag:yaml.org,2002:str', data) |
| 146 | |
| 147 | def represent_binary(self, data): |
| 148 | if hasattr(base64, 'encodebytes'): |
| 149 | data = base64.encodebytes(data).decode('ascii') |
| 150 | else: |
| 151 | data = base64.encodestring(data).decode('ascii') |
| 152 | return self.represent_scalar('tag:yaml.org,2002:binary', data, style='|') |
| 153 | |
| 154 | def represent_bool(self, data): |
| 155 | if data: |
| 156 | value = 'true' |
| 157 | else: |
| 158 | value = 'false' |
| 159 | return self.represent_scalar('tag:yaml.org,2002:bool', value) |
| 160 | |
| 161 | def represent_int(self, data): |
| 162 | return self.represent_scalar('tag:yaml.org,2002:int', str(data)) |
| 163 | |
| 164 | inf_value = 1e300 |
| 165 | while repr(inf_value) != repr(inf_value*inf_value): |
| 166 | inf_value *= inf_value |
| 167 | |
| 168 | def represent_float(self, data): |
| 169 | if data != data or (data == 0.0 and data == 1.0): |
| 170 | value = '.nan' |
| 171 | elif data == self.inf_value: |
| 172 | value = '.inf' |
| 173 | elif data == -self.inf_value: |
| 174 | value = '-.inf' |
| 175 | else: |
| 176 | value = repr(data).lower() |
| 177 | # Note that in some cases `repr(data)` represents a float number |
| 178 | # without the decimal parts. For instance: |
| 179 | # >>> repr(1e17) |
| 180 | # '1e17' |
| 181 | # Unfortunately, this is not a valid float representation according |
| 182 | # to the definition of the `!!float` tag. We fix this by adding |
| 183 | # '.0' before the 'e' symbol. |
| 184 | if '.' not in value and 'e' in value: |
| 185 | value = value.replace('e', '.0e', 1) |
| 186 | return self.represent_scalar('tag:yaml.org,2002:float', value) |
| 187 | |
| 188 | def represent_list(self, data): |
| 189 | #pairs = (len(data) > 0 and isinstance(data, list)) |
| 190 | #if pairs: |
| 191 | # for item in data: |
| 192 | # if not isinstance(item, tuple) or len(item) != 2: |
| 193 | # pairs = False |
| 194 | # break |
| 195 | #if not pairs: |
| 196 | return self.represent_sequence('tag:yaml.org,2002:seq', data) |
| 197 | #value = [] |
| 198 | #for item_key, item_value in data: |
| 199 | # value.append(self.represent_mapping(u'tag:yaml.org,2002:map', |
| 200 | # [(item_key, item_value)])) |
| 201 | #return SequenceNode(u'tag:yaml.org,2002:pairs', value) |
| 202 | |
| 203 | def represent_dict(self, data): |
| 204 | return self.represent_mapping('tag:yaml.org,2002:map', data) |
| 205 | |
| 206 | def represent_set(self, data): |
| 207 | value = {} |
| 208 | for key in data: |
| 209 | value[key] = None |
| 210 | return self.represent_mapping('tag:yaml.org,2002:set', value) |
| 211 | |
| 212 | def represent_date(self, data): |
| 213 | value = data.isoformat() |
| 214 | return self.represent_scalar('tag:yaml.org,2002:timestamp', value) |
| 215 | |
| 216 | def represent_datetime(self, data): |
| 217 | value = data.isoformat(' ') |
| 218 | return self.represent_scalar('tag:yaml.org,2002:timestamp', value) |
| 219 | |
| 220 | def represent_yaml_object(self, tag, data, cls, flow_style=None): |
| 221 | if hasattr(data, '__getstate__'): |
| 222 | state = data.__getstate__() |
| 223 | else: |
| 224 | state = data.__dict__.copy() |
| 225 | return self.represent_mapping(tag, state, flow_style=flow_style) |
| 226 | |
| 227 | def represent_undefined(self, data): |
| 228 | raise RepresenterError("cannot represent an object: %s" % data) |
| 229 | |
| 230 | SafeRepresenter.add_representer(type(None), |
| 231 | SafeRepresenter.represent_none) |
| 232 | |
| 233 | SafeRepresenter.add_representer(str, |
| 234 | SafeRepresenter.represent_str) |
| 235 | |
| 236 | SafeRepresenter.add_representer(bytes, |
| 237 | SafeRepresenter.represent_binary) |
| 238 | |
| 239 | SafeRepresenter.add_representer(bool, |
| 240 | SafeRepresenter.represent_bool) |
| 241 | |
| 242 | SafeRepresenter.add_representer(int, |
| 243 | SafeRepresenter.represent_int) |
| 244 | |
| 245 | SafeRepresenter.add_representer(float, |
| 246 | SafeRepresenter.represent_float) |
| 247 | |
| 248 | SafeRepresenter.add_representer(list, |
| 249 | SafeRepresenter.represent_list) |
| 250 | |
| 251 | SafeRepresenter.add_representer(tuple, |
| 252 | SafeRepresenter.represent_list) |
| 253 | |
| 254 | SafeRepresenter.add_representer(dict, |
| 255 | SafeRepresenter.represent_dict) |
| 256 | |
| 257 | SafeRepresenter.add_representer(set, |
| 258 | SafeRepresenter.represent_set) |
| 259 | |
| 260 | SafeRepresenter.add_representer(datetime.date, |
| 261 | SafeRepresenter.represent_date) |
| 262 | |
| 263 | SafeRepresenter.add_representer(datetime.datetime, |
| 264 | SafeRepresenter.represent_datetime) |
| 265 | |
| 266 | SafeRepresenter.add_representer(None, |
| 267 | SafeRepresenter.represent_undefined) |
| 268 | |
| 269 | class Representer(SafeRepresenter): |
| 270 | |
| 271 | def represent_complex(self, data): |
| 272 | if data.imag == 0.0: |
| 273 | data = '%r' % data.real |
| 274 | elif data.real == 0.0: |
| 275 | data = '%rj' % data.imag |
| 276 | elif data.imag > 0: |
| 277 | data = '%r+%rj' % (data.real, data.imag) |
| 278 | else: |
| 279 | data = '%r%rj' % (data.real, data.imag) |
| 280 | return self.represent_scalar('tag:yaml.org,2002:python/complex', data) |
| 281 | |
| 282 | def represent_tuple(self, data): |
| 283 | return self.represent_sequence('tag:yaml.org,2002:python/tuple', data) |
| 284 | |
| 285 | def represent_name(self, data): |
| 286 | name = '%s.%s' % (data.__module__, data.__name__) |
| 287 | return self.represent_scalar('tag:yaml.org,2002:python/name:'+name, '') |
| 288 | |
| 289 | def represent_module(self, data): |
| 290 | return self.represent_scalar( |
| 291 | 'tag:yaml.org,2002:python/module:'+data.__name__, '') |
| 292 | |
| 293 | def represent_object(self, data): |
| 294 | # We use __reduce__ API to save the data. data.__reduce__ returns |
| 295 | # a tuple of length 2-5: |
| 296 | # (function, args, state, listitems, dictitems) |
| 297 | |
| 298 | # For reconstructing, we calls function(*args), then set its state, |
| 299 | # listitems, and dictitems if they are not None. |
| 300 | |
| 301 | # A special case is when function.__name__ == '__newobj__'. In this |
| 302 | # case we create the object with args[0].__new__(*args). |
| 303 | |
| 304 | # Another special case is when __reduce__ returns a string - we don't |
| 305 | # support it. |
| 306 | |
| 307 | # We produce a !!python/object, !!python/object/new or |
| 308 | # !!python/object/apply node. |
| 309 | |
| 310 | cls = type(data) |
| 311 | if cls in copyreg.dispatch_table: |
| 312 | reduce = copyreg.dispatch_table[cls](data) |
| 313 | elif hasattr(data, '__reduce_ex__'): |
| 314 | reduce = data.__reduce_ex__(2) |
| 315 | elif hasattr(data, '__reduce__'): |
| 316 | reduce = data.__reduce__() |
| 317 | else: |
| 318 | raise RepresenterError("cannot represent object: %r" % data) |
| 319 | reduce = (list(reduce)+[None]*5)[:5] |
| 320 | function, args, state, listitems, dictitems = reduce |
| 321 | args = list(args) |
| 322 | if state is None: |
| 323 | state = {} |
| 324 | if listitems is not None: |
| 325 | listitems = list(listitems) |
| 326 | if dictitems is not None: |
| 327 | dictitems = dict(dictitems) |
| 328 | if function.__name__ == '__newobj__': |
| 329 | function = args[0] |
| 330 | args = args[1:] |
| 331 | tag = 'tag:yaml.org,2002:python/object/new:' |
| 332 | newobj = True |
| 333 | else: |
| 334 | tag = 'tag:yaml.org,2002:python/object/apply:' |
| 335 | newobj = False |
| 336 | function_name = '%s.%s' % (function.__module__, function.__name__) |
| 337 | if not args and not listitems and not dictitems \ |
| 338 | and isinstance(state, dict) and newobj: |
| 339 | return self.represent_mapping( |
| 340 | 'tag:yaml.org,2002:python/object:'+function_name, state) |
| 341 | if not listitems and not dictitems \ |
| 342 | and isinstance(state, dict) and not state: |
| 343 | return self.represent_sequence(tag+function_name, args) |
| 344 | value = {} |
| 345 | if args: |
| 346 | value['args'] = args |
| 347 | if state or not isinstance(state, dict): |
| 348 | value['state'] = state |
| 349 | if listitems: |
| 350 | value['listitems'] = listitems |
| 351 | if dictitems: |
| 352 | value['dictitems'] = dictitems |
| 353 | return self.represent_mapping(tag+function_name, value) |
| 354 | |
| 355 | Representer.add_representer(complex, |
| 356 | Representer.represent_complex) |
| 357 | |
| 358 | Representer.add_representer(tuple, |
| 359 | Representer.represent_tuple) |
| 360 | |
| 361 | Representer.add_representer(type, |
| 362 | Representer.represent_name) |
| 363 | |
| 364 | Representer.add_representer(types.FunctionType, |
| 365 | Representer.represent_name) |
| 366 | |
| 367 | Representer.add_representer(types.BuiltinFunctionType, |
| 368 | Representer.represent_name) |
| 369 | |
| 370 | Representer.add_representer(types.ModuleType, |
| 371 | Representer.represent_module) |
| 372 | |
| 373 | Representer.add_multi_representer(object, |
| 374 | Representer.represent_object) |
| 375 |