1
2__all__ = ['BaseRepresenter', 'SafeRepresenter', 'Representer',
3    'RepresenterError']
4
5from .error import *
6from .nodes import *
7
8import datetime, copyreg, types, base64, collections
9
10class RepresenterError(YAMLError):
11    pass
12
13class BaseRepresenter:
14
15    yaml_representers = {}
16    yaml_multi_representers = {}
17
18    def __init__(self, default_style=None, default_flow_style=False, sort_keys=True):
19        self.default_style = default_style
20        self.sort_keys = sort_keys
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            if self.sort_keys:
112                try:
113                    mapping = sorted(mapping)
114                except TypeError:
115                    pass
116        for item_key, item_value in mapping:
117            node_key = self.represent_data(item_key)
118            node_value = self.represent_data(item_value)
119            if not (isinstance(node_key, ScalarNode) and not node_key.style):
120                best_style = False
121            if not (isinstance(node_value, ScalarNode) and not node_value.style):
122                best_style = False
123            value.append((node_key, node_value))
124        if flow_style is None:
125            if self.default_flow_style is not None:
126                node.flow_style = self.default_flow_style
127            else:
128                node.flow_style = best_style
129        return node
130
131    def ignore_aliases(self, data):
132        return False
133
134class SafeRepresenter(BaseRepresenter):
135
136    def ignore_aliases(self, data):
137        if data is None:
138            return True
139        if isinstance(data, tuple) and data == ():
140            return True
141        if isinstance(data, (str, bytes, bool, int, float)):
142            return True
143
144    def represent_none(self, data):
145        return self.represent_scalar('tag:yaml.org,2002:null', 'null')
146
147    def represent_str(self, data):
148        return self.represent_scalar('tag:yaml.org,2002:str', data)
149
150    def represent_binary(self, data):
151        if hasattr(base64, 'encodebytes'):
152            data = base64.encodebytes(data).decode('ascii')
153        else:
154            data = base64.encodestring(data).decode('ascii')
155        return self.represent_scalar('tag:yaml.org,2002:binary', data, style='|')
156
157    def represent_bool(self, data):
158        if data:
159            value = 'true'
160        else:
161            value = 'false'
162        return self.represent_scalar('tag:yaml.org,2002:bool', value)
163
164    def represent_int(self, data):
165        return self.represent_scalar('tag:yaml.org,2002:int', str(data))
166
167    inf_value = 1e300
168    while repr(inf_value) != repr(inf_value*inf_value):
169        inf_value *= inf_value
170
171    def represent_float(self, data):
172        if data != data or (data == 0.0 and data == 1.0):
173            value = '.nan'
174        elif data == self.inf_value:
175            value = '.inf'
176        elif data == -self.inf_value:
177            value = '-.inf'
178        else:
179            value = repr(data).lower()
180            # Note that in some cases `repr(data)` represents a float number
181            # without the decimal parts.  For instance:
182            #   >>> repr(1e17)
183            #   '1e17'
184            # Unfortunately, this is not a valid float representation according
185            # to the definition of the `!!float` tag.  We fix this by adding
186            # '.0' before the 'e' symbol.
187            if '.' not in value and 'e' in value:
188                value = value.replace('e', '.0e', 1)
189        return self.represent_scalar('tag:yaml.org,2002:float', value)
190
191    def represent_list(self, data):
192        #pairs = (len(data) > 0 and isinstance(data, list))
193        #if pairs:
194        #    for item in data:
195        #        if not isinstance(item, tuple) or len(item) != 2:
196        #            pairs = False
197        #            break
198        #if not pairs:
199            return self.represent_sequence('tag:yaml.org,2002:seq', data)
200        #value = []
201        #for item_key, item_value in data:
202        #    value.append(self.represent_mapping(u'tag:yaml.org,2002:map',
203        #        [(item_key, item_value)]))
204        #return SequenceNode(u'tag:yaml.org,2002:pairs', value)
205
206    def represent_dict(self, data):
207        return self.represent_mapping('tag:yaml.org,2002:map', data)
208
209    def represent_set(self, data):
210        value = {}
211        for key in data:
212            value[key] = None
213        return self.represent_mapping('tag:yaml.org,2002:set', value)
214
215    def represent_date(self, data):
216        value = data.isoformat()
217        return self.represent_scalar('tag:yaml.org,2002:timestamp', value)
218
219    def represent_datetime(self, data):
220        value = data.isoformat(' ')
221        return self.represent_scalar('tag:yaml.org,2002:timestamp', value)
222
223    def represent_yaml_object(self, tag, data, cls, flow_style=None):
224        if hasattr(data, '__getstate__'):
225            state = data.__getstate__()
226        else:
227            state = data.__dict__.copy()
228        return self.represent_mapping(tag, state, flow_style=flow_style)
229
230    def represent_undefined(self, data):
231        raise RepresenterError("cannot represent an object", data)
232
233SafeRepresenter.add_representer(type(None),
234        SafeRepresenter.represent_none)
235
236SafeRepresenter.add_representer(str,
237        SafeRepresenter.represent_str)
238
239SafeRepresenter.add_representer(bytes,
240        SafeRepresenter.represent_binary)
241
242SafeRepresenter.add_representer(bool,
243        SafeRepresenter.represent_bool)
244
245SafeRepresenter.add_representer(int,
246        SafeRepresenter.represent_int)
247
248SafeRepresenter.add_representer(float,
249        SafeRepresenter.represent_float)
250
251SafeRepresenter.add_representer(list,
252        SafeRepresenter.represent_list)
253
254SafeRepresenter.add_representer(tuple,
255        SafeRepresenter.represent_list)
256
257SafeRepresenter.add_representer(dict,
258        SafeRepresenter.represent_dict)
259
260SafeRepresenter.add_representer(set,
261        SafeRepresenter.represent_set)
262
263SafeRepresenter.add_representer(datetime.date,
264        SafeRepresenter.represent_date)
265
266SafeRepresenter.add_representer(datetime.datetime,
267        SafeRepresenter.represent_datetime)
268
269SafeRepresenter.add_representer(None,
270        SafeRepresenter.represent_undefined)
271
272class Representer(SafeRepresenter):
273
274    def represent_complex(self, data):
275        if data.imag == 0.0:
276            data = '%r' % data.real
277        elif data.real == 0.0:
278            data = '%rj' % data.imag
279        elif data.imag > 0:
280            data = '%r+%rj' % (data.real, data.imag)
281        else:
282            data = '%r%rj' % (data.real, data.imag)
283        return self.represent_scalar('tag:yaml.org,2002:python/complex', data)
284
285    def represent_tuple(self, data):
286        return self.represent_sequence('tag:yaml.org,2002:python/tuple', data)
287
288    def represent_name(self, data):
289        name = '%s.%s' % (data.__module__, data.__name__)
290        return self.represent_scalar('tag:yaml.org,2002:python/name:'+name, '')
291
292    def represent_module(self, data):
293        return self.represent_scalar(
294                'tag:yaml.org,2002:python/module:'+data.__name__, '')
295
296    def represent_object(self, data):
297        # We use __reduce__ API to save the data. data.__reduce__ returns
298        # a tuple of length 2-5:
299        #   (function, args, state, listitems, dictitems)
300
301        # For reconstructing, we calls function(*args), then set its state,
302        # listitems, and dictitems if they are not None.
303
304        # A special case is when function.__name__ == '__newobj__'. In this
305        # case we create the object with args[0].__new__(*args).
306
307        # Another special case is when __reduce__ returns a string - we don't
308        # support it.
309
310        # We produce a !!python/object, !!python/object/new or
311        # !!python/object/apply node.
312
313        cls = type(data)
314        if cls in copyreg.dispatch_table:
315            reduce = copyreg.dispatch_table[cls](data)
316        elif hasattr(data, '__reduce_ex__'):
317            reduce = data.__reduce_ex__(2)
318        elif hasattr(data, '__reduce__'):
319            reduce = data.__reduce__()
320        else:
321            raise RepresenterError("cannot represent an object", data)
322        reduce = (list(reduce)+[None]*5)[:5]
323        function, args, state, listitems, dictitems = reduce
324        args = list(args)
325        if state is None:
326            state = {}
327        if listitems is not None:
328            listitems = list(listitems)
329        if dictitems is not None:
330            dictitems = dict(dictitems)
331        if function.__name__ == '__newobj__':
332            function = args[0]
333            args = args[1:]
334            tag = 'tag:yaml.org,2002:python/object/new:'
335            newobj = True
336        else:
337            tag = 'tag:yaml.org,2002:python/object/apply:'
338            newobj = False
339        function_name = '%s.%s' % (function.__module__, function.__name__)
340        if not args and not listitems and not dictitems \
341                and isinstance(state, dict) and newobj:
342            return self.represent_mapping(
343                    'tag:yaml.org,2002:python/object:'+function_name, state)
344        if not listitems and not dictitems  \
345                and isinstance(state, dict) and not state:
346            return self.represent_sequence(tag+function_name, args)
347        value = {}
348        if args:
349            value['args'] = args
350        if state or not isinstance(state, dict):
351            value['state'] = state
352        if listitems:
353            value['listitems'] = listitems
354        if dictitems:
355            value['dictitems'] = dictitems
356        return self.represent_mapping(tag+function_name, value)
357
358    def represent_ordered_dict(self, data):
359        # Provide uniform representation across different Python versions.
360        data_type = type(data)
361        tag = 'tag:yaml.org,2002:python/object/apply:%s.%s' \
362                % (data_type.__module__, data_type.__name__)
363        items = [[key, value] for key, value in data.items()]
364        return self.represent_sequence(tag, [items])
365
366Representer.add_representer(complex,
367        Representer.represent_complex)
368
369Representer.add_representer(tuple,
370        Representer.represent_tuple)
371
372Representer.add_representer(type,
373        Representer.represent_name)
374
375Representer.add_representer(collections.OrderedDict,
376        Representer.represent_ordered_dict)
377
378Representer.add_representer(types.FunctionType,
379        Representer.represent_name)
380
381Representer.add_representer(types.BuiltinFunctionType,
382        Representer.represent_name)
383
384Representer.add_representer(types.ModuleType,
385        Representer.represent_module)
386
387Representer.add_multi_representer(object,
388        Representer.represent_object)
389
390