Ticket #5: constructor.py

File constructor.py, 16.3 KB (added by pkmurphy@…, 6 years ago)

Patch on file for problem

Line 
1
2__all__ = ['BaseConstructor', 'SafeConstructor', 'Constructor',
3    'ConstructorError']
4
5# This contains revisions to allow recursive loads in this file.
6# All modifications are marked by PKM2006
7
8
9from error import *
10from nodes import *
11from composer import *
12
13try:
14    import datetime
15    datetime_available = True
16except ImportError:
17    datetime_available = False
18
19try:
20    set
21except NameError:
22    from sets import Set as set
23
24import binascii, re
25
26class ConstructorError(MarkedYAMLError):
27    pass
28
29class BaseConstructor(Composer):
30
31    yaml_constructors = {}
32    yaml_multi_constructors = {}
33
34    def __init__(self):
35        self.constructed_objects = {}
36
37    def check_data(self):
38        # If there are more documents available?
39        return self.check_node()
40
41    def get_data(self):
42        # Construct and return the next document.
43        if self.check_node():
44            return self.construct_document(self.get_node())
45
46    def __iter__(self):
47        # Iterator protocol.
48        while self.check_node():
49            yield self.construct_document(self.get_node())
50
51    def construct_document(self, node):
52        data = self.construct_object(node)
53        self.constructed_objects = {}
54        return data
55
56    def construct_object(self, node):
57        if node in self.constructed_objects:
58            return self.constructed_objects[node]
59        constructor = None
60        if node.tag in self.yaml_constructors:
61            constructor = lambda node: self.yaml_constructors[node.tag](self, node)
62        else:
63            for tag_prefix in self.yaml_multi_constructors:
64                if node.tag.startswith(tag_prefix):
65                    tag_suffix = node.tag[len(tag_prefix):]
66                    constructor = lambda node:  \
67                            self.yaml_multi_constructors[tag_prefix](self, tag_suffix, node)
68                break
69            else:
70                if None in self.yaml_multi_constructors:
71                    constructor = lambda node:  \
72                            self.yaml_multi_constructors[None](self, node.tag, node)
73                elif None in self.yaml_constructors:
74                    constructor = lambda node:  \
75                            self.yaml_constructors[None](self, node)
76                elif isinstance(node, ScalarNode):
77                    constructor = self.construct_scalar
78                elif isinstance(node, SequenceNode):
79                    constructor = self.construct_sequence
80                elif isinstance(node, MappingNode):
81                    constructor = self.construct_mapping
82        data = constructor(node)
83        self.constructed_objects[node] = data
84        return data
85
86    def construct_scalar(self, node):
87        if not isinstance(node, ScalarNode):
88            if isinstance(node, MappingNode):
89                for key_node in node.value:
90                    if key_node.tag == u'tag:yaml.org,2002:value':
91                        data = self.construct_scalar(node.value[key_node]); # PKM2006
92                        self.constructed_objects[node] = data # PKM2006
93                        return data # PKM2006
94            raise ConstructorError(None, None,
95                    "expected a scalar node, but found %s" % node.id,
96                    node.start_mark)
97        data = node.value; # PKM2006
98        self.constructed_objects[node] = data; # PKM2006
99        return data; # PKM2006
100
101    def construct_sequence(self, node):
102        if not isinstance(node, SequenceNode):
103            raise ConstructorError(None, None,
104                    "expected a sequence node, but found %s" % node.id,
105                    node.start_mark)
106# PKM2006: return [self.construct_object(child) for child in node.value]
107        data = [];
108        self.constructed_objects[node] = data; # PKM2006
109        for child in node.value:
110            data.append(self.construct_object(child));
111        return data;
112
113    def construct_mapping(self, node):
114        if not isinstance(node, MappingNode):
115            raise ConstructorError(None, None,
116                    "expected a mapping node, but found %s" % node.id,
117                    node.start_mark)
118        mapping = {}
119        self.constructed_objects[node] = mapping; # PKM2006       
120        merge = None
121        for key_node in node.value:
122            if key_node.tag == u'tag:yaml.org,2002:merge':
123                if merge is not None:
124                    raise ConstructorError("while constructing a mapping", node.start_mark,
125                            "found duplicate merge key", key_node.start_mark)
126                value_node = node.value[key_node]
127                if isinstance(value_node, MappingNode):
128                    merge = [self.construct_mapping(value_node)]
129                elif isinstance(value_node, SequenceNode):
130                    merge = []
131                    for subnode in value_node.value:
132                        if not isinstance(subnode, MappingNode):
133                            raise ConstructorError("while constructing a mapping",
134                                    node.start_mark,
135                                    "expected a mapping for merging, but found %s"
136                                    % subnode.id, subnode.start_mark)
137                        merge.append(self.construct_mapping(subnode))
138                    merge.reverse()
139                else:
140                    raise ConstructorError("while constructing a mapping", node.start_mark,
141                            "expected a mapping or list of mappings for merging, but found %s"
142                            % value_node.id, value_node.start_mark)
143            elif key_node.tag == u'tag:yaml.org,2002:value':
144                if '=' in mapping:
145                    raise ConstructorError("while construction a mapping", node.start_mark,
146                            "found duplicate value key", key_node.start_mark)
147                value = self.construct_object(node.value[key_node])
148                mapping['='] = value
149            else:
150                key = self.construct_object(key_node)
151                try:
152                    duplicate_key = key in mapping
153                except TypeError, exc:
154                    raise ConstructorError("while constructing a mapping", node.start_mark,
155                            "found unacceptable key (%s)" % exc, key_node.start_mark)
156                if duplicate_key:
157                    raise ConstructorError("while constructing a mapping", node.start_mark,
158                            "found duplicate key", key_node.start_mark)
159                value = self.construct_object(node.value[key_node])
160                mapping[key] = value
161        if merge is not None:
162            merge.append(mapping)
163            mapping.clear(); # PKM - better than mapping = {} as we need to connect one to another.
164            for submapping in merge:
165                mapping.update(submapping)
166        return mapping
167
168    def construct_pairs(self, node):
169        if not isinstance(node, MappingNode):
170            raise ConstructorError(None, None,
171                    "expected a mapping node, but found %s" % node.id,
172                    node.start_mark)
173        pairs = []
174        self.constructed_objects[node] = data; # PKM2006
175        for key_node in node.value:
176            key = self.construct_object(key_node)
177            value = self.construct_object(node.value[key_node])
178            pairs.append((key, value))
179        return pairs
180
181    def add_constructor(cls, tag, constructor):
182        if not 'yaml_constructors' in cls.__dict__:
183            cls.yaml_constructors = cls.yaml_constructors.copy()
184        cls.yaml_constructors[tag] = constructor
185    add_constructor = classmethod(add_constructor)
186
187    def add_multi_constructor(cls, tag_prefix, multi_constructor):
188        if not 'yaml_multi_constructors' in cls.__dict__:
189            cls.yaml_multi_constructors = cls.yaml_multi_constructors.copy()
190        cls.yaml_multi_constructors[tag_prefix] = multi_constructor
191    add_multi_constructor = classmethod(add_multi_constructor)
192
193class SafeConstructor(BaseConstructor):
194
195    def construct_yaml_null(self, node):
196        self.construct_scalar(node)
197        return None
198
199    bool_values = {
200        u'yes':     True,
201        u'no':      False,
202        u'true':    True,
203        u'false':   False,
204        u'on':      True,
205        u'off':     False,
206    }
207
208    def construct_yaml_bool(self, node):
209        value = self.construct_scalar(node)
210        return self.bool_values[value.lower()]
211
212    def construct_yaml_int(self, node):
213        value = str(self.construct_scalar(node))
214        value = value.replace('_', '')
215        sign = +1
216        if value[0] == '-':
217            sign = -1
218        if value[0] in '+-':
219            value = value[1:]
220        if value == '0':
221            return 0
222        elif value.startswith('0b'):
223            return sign*int(value[2:], 2)
224        elif value.startswith('0x'):
225            return sign*int(value[2:], 16)
226        elif value[0] == '0':
227            return sign*int(value, 8)
228        elif ':' in value:
229            digits = [int(part) for part in value.split(':')]
230            digits.reverse()
231            base = 1
232            value = 0
233            for digit in digits:
234                value += digit*base
235                base *= 60
236            return sign*value
237        else:
238            return sign*int(value)
239
240    inf_value = 1e300000
241    nan_value = inf_value/inf_value
242
243    def construct_yaml_float(self, node):
244        value = str(self.construct_scalar(node))
245        value = value.replace('_', '')
246        sign = +1
247        if value[0] == '-':
248            sign = -1
249        if value[0] in '+-':
250            value = value[1:]
251        if value.lower() == '.inf':
252            return sign*self.inf_value
253        elif value.lower() == '.nan':
254            return self.nan_value
255        elif ':' in value:
256            digits = [float(part) for part in value.split(':')]
257            digits.reverse()
258            base = 1
259            value = 0.0
260            for digit in digits:
261                value += digit*base
262                base *= 60
263            return sign*value
264        else:
265            return float(value)
266
267    def construct_yaml_binary(self, node):
268        value = self.construct_scalar(node)
269        try:
270            return str(value).decode('base64')
271        except (binascii.Error, UnicodeEncodeError), exc:
272            raise ConstructorError(None, None,
273                    "failed to decode base64 data: %s" % exc, node.start_mark) 
274
275    timestamp_regexp = re.compile(
276            ur'''^(?P<year>[0-9][0-9][0-9][0-9])
277                -(?P<month>[0-9][0-9]?)
278                -(?P<day>[0-9][0-9]?)
279                (?:(?:[Tt]|[ \t]+)
280                (?P<hour>[0-9][0-9]?)
281                :(?P<minute>[0-9][0-9])
282                :(?P<second>[0-9][0-9])
283                (?:\.(?P<fraction>[0-9]*))?
284                (?:[ \t]*(?:Z|(?P<tz_hour>[-+][0-9][0-9]?)
285                (?::(?P<tz_minute>[0-9][0-9])?)?))?)?$''', re.X)
286
287    def construct_yaml_timestamp(self, node):
288        value = self.construct_scalar(node)
289        match = self.timestamp_regexp.match(node.value)
290        values = match.groupdict()
291        for key in values:
292            if values[key]:
293                values[key] = int(values[key])
294            else:
295                values[key] = 0
296        fraction = values['fraction']
297        if fraction:
298            while 10*fraction < 1000000:
299                fraction *= 10
300            values['fraction'] = fraction
301        stamp = datetime.datetime(values['year'], values['month'], values['day'],
302                values['hour'], values['minute'], values['second'], values['fraction'])
303        diff = datetime.timedelta(hours=values['tz_hour'], minutes=values['tz_minute'])
304        return stamp-diff
305
306    def construct_yaml_omap(self, node):
307        # Note: we do not check for duplicate keys, because it's too
308        # CPU-expensive.
309        if not isinstance(node, SequenceNode):
310            raise ConstructorError("while constructing an ordered map", node.start_mark,
311                    "expected a sequence, but found %s" % node.id, node.start_mark)
312        omap = []
313        for subnode in node.value:
314            if not isinstance(subnode, MappingNode):
315                raise ConstructorError("while constructing an ordered map", node.start_mark,
316                        "expected a mapping of length 1, but found %s" % subnode.id,
317                        subnode.start_mark)
318            if len(subnode.value) != 1:
319                raise ConstructorError("while constructing an ordered map", node.start_mark,
320                        "expected a single mapping item, but found %d items" % len(subnode.value),
321                        subnode.start_mark)
322            key_node = subnode.value.keys()[0]
323            key = self.construct_object(key_node)
324            value = self.construct_object(subnode.value[key_node])
325            omap.append((key, value))
326        return omap
327
328    def construct_yaml_pairs(self, node):
329        # Note: the same code as `construct_yaml_omap`.
330        if not isinstance(node, SequenceNode):
331            raise ConstructorError("while constructing pairs", node.start_mark,
332                    "expected a sequence, but found %s" % node.id, node.start_mark)
333        pairs = []
334        for subnode in node.value:
335            if not isinstance(subnode, MappingNode):
336                raise ConstructorError("while constructing pairs", node.start_mark,
337                        "expected a mapping of length 1, but found %s" % subnode.id,
338                        subnode.start_mark)
339            if len(subnode.value) != 1:
340                raise ConstructorError("while constructing pairs", node.start_mark,
341                        "expected a single mapping item, but found %d items" % len(subnode.value),
342                        subnode.start_mark)
343            key_node = subnode.value.keys()[0]
344            key = self.construct_object(key_node)
345            value = self.construct_object(subnode.value[key_node])
346            pairs.append((key, value))
347        return pairs
348
349    def construct_yaml_set(self, node):
350        value = self.construct_mapping(node)
351        return set(value)
352
353    def construct_yaml_str(self, node):
354        value = self.construct_scalar(node)
355        try:
356            return str(value)
357        except UnicodeEncodeError:
358            return value
359
360    def construct_yaml_seq(self, node):
361        return self.construct_sequence(node)
362
363    def construct_yaml_map(self, node):
364        return self.construct_mapping(node)
365
366    def construct_yaml_object(self, node, cls):
367        mapping = self.construct_mapping(node)
368        state = {}
369        for key in mapping:
370            state[key.replace('-', '_')] = mapping[key]
371        data = cls.__new__(cls)
372        if hasattr(data, '__setstate__'):
373            data.__setstate__(mapping)
374        else:
375            data.__dict__.update(mapping)
376        return data
377
378    def construct_undefined(self, node):
379        raise ConstructorError(None, None,
380                "could not determine a constructor for the tag %r" % node.tag.encode('utf-8'),
381                node.start_mark)
382
383SafeConstructor.add_constructor(
384        u'tag:yaml.org,2002:null',
385        SafeConstructor.construct_yaml_null)
386
387SafeConstructor.add_constructor(
388        u'tag:yaml.org,2002:bool',
389        SafeConstructor.construct_yaml_bool)
390
391SafeConstructor.add_constructor(
392        u'tag:yaml.org,2002:int',
393        SafeConstructor.construct_yaml_int)
394
395SafeConstructor.add_constructor(
396        u'tag:yaml.org,2002:float',
397        SafeConstructor.construct_yaml_float)
398
399SafeConstructor.add_constructor(
400        u'tag:yaml.org,2002:binary',
401        SafeConstructor.construct_yaml_binary)
402
403if datetime_available:
404    SafeConstructor.add_constructor(
405            u'tag:yaml.org,2002:timestamp',
406            SafeConstructor.construct_yaml_timestamp)
407
408SafeConstructor.add_constructor(
409        u'tag:yaml.org,2002:omap',
410        SafeConstructor.construct_yaml_omap)
411
412SafeConstructor.add_constructor(
413        u'tag:yaml.org,2002:pairs',
414        SafeConstructor.construct_yaml_pairs)
415
416SafeConstructor.add_constructor(
417        u'tag:yaml.org,2002:set',
418        SafeConstructor.construct_yaml_set)
419
420SafeConstructor.add_constructor(
421        u'tag:yaml.org,2002:str',
422        SafeConstructor.construct_yaml_str)
423
424SafeConstructor.add_constructor(
425        u'tag:yaml.org,2002:seq',
426        SafeConstructor.construct_yaml_seq)
427
428SafeConstructor.add_constructor(
429        u'tag:yaml.org,2002:map',
430        SafeConstructor.construct_yaml_map)
431
432SafeConstructor.add_constructor(None,
433        SafeConstructor.construct_undefined)
434
435class Constructor(SafeConstructor):
436    pass
437