source: pyyaml/trunk/lib/yaml/constructor.py @ 258

Revision 258, 24.7 KB checked in by xi, 7 years ago (diff)

Make compose() and load() ensure that the input stream contains a single document. Fixes #54.

Line 
1
2__all__ = ['BaseConstructor', 'SafeConstructor', 'Constructor',
3    'ConstructorError']
4
5from error import *
6from nodes import *
7
8import datetime
9
10try:
11    set
12except NameError:
13    from sets import Set as set
14
15import binascii, re, sys, types
16
17class ConstructorError(MarkedYAMLError):
18    pass
19
20class BaseConstructor(object):
21
22    yaml_constructors = {}
23    yaml_multi_constructors = {}
24
25    def __init__(self):
26        self.constructed_objects = {}
27        self.recursive_objects = {}
28        self.state_generators = []
29        self.deep_construct = False
30
31    def check_data(self):
32        # If there are more documents available?
33        return self.check_node()
34
35    def get_data(self):
36        # Construct and return the next document.
37        if self.check_node():
38            return self.construct_document(self.get_node())
39
40    def get_single_data(self):
41        # Ensure that the stream contains a single document and construct it.
42        node = self.get_single_node()
43        if node is not None:
44            return self.construct_document(node)
45        return None
46
47    def construct_document(self, node):
48        data = self.construct_object(node)
49        while self.state_generators:
50            state_generators = self.state_generators
51            self.state_generators = []
52            for generator in state_generators:
53                for dummy in generator:
54                    pass
55        self.constructed_objects = {}
56        self.recursive_objects = {}
57        self.deep_construct = False
58        return data
59
60    def construct_object(self, node, deep=False):
61        if deep:
62            old_deep = self.deep_construct
63            self.deep_construct = True
64        if node in self.constructed_objects:
65            return self.constructed_objects[node]
66        if node in self.recursive_objects:
67            raise ConstructorError(None, None,
68                    "found unconstructable recursive node", node.start_mark)
69        self.recursive_objects[node] = None
70        constructor = None
71        state_constructor = None
72        tag_suffix = None
73        if node.tag in self.yaml_constructors:
74            constructor = self.yaml_constructors[node.tag]
75        else:
76            for tag_prefix in self.yaml_multi_constructors:
77                if node.tag.startswith(tag_prefix):
78                    tag_suffix = node.tag[len(tag_prefix):]
79                    constructor = self.yaml_multi_constructors[tag_prefix]
80                    break
81            else:
82                if None in self.yaml_multi_constructors:
83                    tag_suffix = node.tag
84                    constructor = self.yaml_multi_constructors[None]
85                elif None in self.yaml_constructors:
86                    constructor = self.yaml_constructors[None]
87                elif isinstance(node, ScalarNode):
88                    constructor = self.__class__.construct_scalar
89                elif isinstance(node, SequenceNode):
90                    constructor = self.__class__.construct_sequence
91                elif isinstance(node, MappingNode):
92                    constructor = self.__class__.construct_mapping
93        if tag_suffix is None:
94            data = constructor(self, node)
95        else:
96            data = constructor(self, tag_suffix, node)
97        if isinstance(data, types.GeneratorType):
98            generator = data
99            data = generator.next()
100            if self.deep_construct:
101                for dummy in generator:
102                    pass
103            else:
104                self.state_generators.append(generator)
105        self.constructed_objects[node] = data
106        del self.recursive_objects[node]
107        if deep:
108            self.deep_construct = old_deep
109        return data
110
111    def construct_scalar(self, node):
112        if not isinstance(node, ScalarNode):
113            raise ConstructorError(None, None,
114                    "expected a scalar node, but found %s" % node.id,
115                    node.start_mark)
116        return node.value
117
118    def construct_sequence(self, node, deep=False):
119        if not isinstance(node, SequenceNode):
120            raise ConstructorError(None, None,
121                    "expected a sequence node, but found %s" % node.id,
122                    node.start_mark)
123        return [self.construct_object(child, deep=deep)
124                for child in node.value]
125
126    def construct_mapping(self, node, deep=False):
127        if not isinstance(node, MappingNode):
128            raise ConstructorError(None, None,
129                    "expected a mapping node, but found %s" % node.id,
130                    node.start_mark)
131        mapping = {}
132        for key_node, value_node in node.value:
133            key = self.construct_object(key_node, deep=deep)
134            try:
135                hash(key)
136            except TypeError, exc:
137                raise ConstructorError("while constructing a mapping", node.start_mark,
138                        "found unacceptable key (%s)" % exc, key_node.start_mark)
139            value = self.construct_object(value_node, deep=deep)
140            mapping[key] = value
141        return mapping
142
143    def construct_pairs(self, node, deep=False):
144        if not isinstance(node, MappingNode):
145            raise ConstructorError(None, None,
146                    "expected a mapping node, but found %s" % node.id,
147                    node.start_mark)
148        pairs = []
149        for key_node, value_node in node.value:
150            key = self.construct_object(key_node, deep=deep)
151            value = self.construct_object(value_node, deep=deep)
152            pairs.append((key, value))
153        return pairs
154
155    def add_constructor(cls, tag, constructor):
156        if not 'yaml_constructors' in cls.__dict__:
157            cls.yaml_constructors = cls.yaml_constructors.copy()
158        cls.yaml_constructors[tag] = constructor
159    add_constructor = classmethod(add_constructor)
160
161    def add_multi_constructor(cls, tag_prefix, multi_constructor):
162        if not 'yaml_multi_constructors' in cls.__dict__:
163            cls.yaml_multi_constructors = cls.yaml_multi_constructors.copy()
164        cls.yaml_multi_constructors[tag_prefix] = multi_constructor
165    add_multi_constructor = classmethod(add_multi_constructor)
166
167class SafeConstructor(BaseConstructor):
168
169    def construct_scalar(self, node):
170        if isinstance(node, MappingNode):
171            for key_node, value_node in node.value:
172                if key_node.tag == u'tag:yaml.org,2002:value':
173                    return self.construct_scalar(value_node)
174        return BaseConstructor.construct_scalar(self, node)
175
176    def flatten_mapping(self, node):
177        merge = []
178        index = 0
179        while index < len(node.value):
180            key_node, value_node = node.value[index]
181            if key_node.tag == u'tag:yaml.org,2002:merge':
182                del node.value[index]
183                if isinstance(value_node, MappingNode):
184                    self.flatten_mapping(value_node)
185                    merge.extend(value_node.value)
186                elif isinstance(value_node, SequenceNode):
187                    submerge = []
188                    for subnode in value_node.value:
189                        if not isinstance(subnode, MappingNode):
190                            raise ConstructorError("while constructing a mapping",
191                                    node.start_mark,
192                                    "expected a mapping for merging, but found %s"
193                                    % subnode.id, subnode.start_mark)
194                        self.flatten_mapping(subnode)
195                        submerge.append(subnode.value)
196                    submerge.reverse()
197                    for value in submerge:
198                        merge.extend(value)
199                else:
200                    raise ConstructorError("while constructing a mapping", node.start_mark,
201                            "expected a mapping or list of mappings for merging, but found %s"
202                            % value_node.id, value_node.start_mark)
203            elif key_node.tag == u'tag:yaml.org,2002:value':
204                key_node.tag = u'tag:yaml.org,2002:str'
205                index += 1
206            else:
207                index += 1
208        if merge:
209            node.value = merge + node.value
210
211    def construct_mapping(self, node, deep=False):
212        if isinstance(node, MappingNode):
213            self.flatten_mapping(node)
214        return BaseConstructor.construct_mapping(self, node, deep=deep)
215
216    def construct_yaml_null(self, node):
217        self.construct_scalar(node)
218        return None
219
220    bool_values = {
221        u'yes':     True,
222        u'no':      False,
223        u'true':    True,
224        u'false':   False,
225        u'on':      True,
226        u'off':     False,
227    }
228
229    def construct_yaml_bool(self, node):
230        value = self.construct_scalar(node)
231        return self.bool_values[value.lower()]
232
233    def construct_yaml_int(self, node):
234        value = str(self.construct_scalar(node))
235        value = value.replace('_', '')
236        sign = +1
237        if value[0] == '-':
238            sign = -1
239        if value[0] in '+-':
240            value = value[1:]
241        if value == '0':
242            return 0
243        elif value.startswith('0b'):
244            return sign*int(value[2:], 2)
245        elif value.startswith('0x'):
246            return sign*int(value[2:], 16)
247        elif value[0] == '0':
248            return sign*int(value, 8)
249        elif ':' in value:
250            digits = [int(part) for part in value.split(':')]
251            digits.reverse()
252            base = 1
253            value = 0
254            for digit in digits:
255                value += digit*base
256                base *= 60
257            return sign*value
258        else:
259            return sign*int(value)
260
261    inf_value = 1e300
262    while inf_value != inf_value*inf_value:
263        inf_value *= inf_value
264    nan_value = -inf_value/inf_value   # Trying to make a quiet NaN (like C99).
265
266    def construct_yaml_float(self, node):
267        value = str(self.construct_scalar(node))
268        value = value.replace('_', '').lower()
269        sign = +1
270        if value[0] == '-':
271            sign = -1
272        if value[0] in '+-':
273            value = value[1:]
274        if value == '.inf':
275            return sign*self.inf_value
276        elif value == '.nan':
277            return self.nan_value
278        elif ':' in value:
279            digits = [float(part) for part in value.split(':')]
280            digits.reverse()
281            base = 1
282            value = 0.0
283            for digit in digits:
284                value += digit*base
285                base *= 60
286            return sign*value
287        else:
288            return sign*float(value)
289
290    def construct_yaml_binary(self, node):
291        value = self.construct_scalar(node)
292        try:
293            return str(value).decode('base64')
294        except (binascii.Error, UnicodeEncodeError), exc:
295            raise ConstructorError(None, None,
296                    "failed to decode base64 data: %s" % exc, node.start_mark) 
297
298    timestamp_regexp = re.compile(
299            ur'''^(?P<year>[0-9][0-9][0-9][0-9])
300                -(?P<month>[0-9][0-9]?)
301                -(?P<day>[0-9][0-9]?)
302                (?:(?:[Tt]|[ \t]+)
303                (?P<hour>[0-9][0-9]?)
304                :(?P<minute>[0-9][0-9])
305                :(?P<second>[0-9][0-9])
306                (?:\.(?P<fraction>[0-9]*))?
307                (?:[ \t]*(?P<tz>Z|(?P<tz_sign>[-+])(?P<tz_hour>[0-9][0-9]?)
308                (?::(?P<tz_minute>[0-9][0-9]))?))?)?$''', re.X)
309
310    def construct_yaml_timestamp(self, node):
311        value = self.construct_scalar(node)
312        match = self.timestamp_regexp.match(node.value)
313        values = match.groupdict()
314        year = int(values['year'])
315        month = int(values['month'])
316        day = int(values['day'])
317        if not values['hour']:
318            return datetime.date(year, month, day)
319        hour = int(values['hour'])
320        minute = int(values['minute'])
321        second = int(values['second'])
322        fraction = 0
323        if values['fraction']:
324            fraction = int(values['fraction'][:6].ljust(6, '0'))
325        delta = None
326        if values['tz_sign']:
327            tz_hour = int(values['tz_hour'])
328            tz_minute = int(values['tz_minute'] or 0)
329            delta = datetime.timedelta(hours=tz_hour, minutes=tz_minute)
330            if values['tz_sign'] == '-':
331                delta = -delta
332        data = datetime.datetime(year, month, day, hour, minute, second, fraction)
333        if delta:
334            data -= delta
335        return data
336
337    def construct_yaml_omap(self, node):
338        # Note: we do not check for duplicate keys, because it's too
339        # CPU-expensive.
340        omap = []
341        yield omap
342        if not isinstance(node, SequenceNode):
343            raise ConstructorError("while constructing an ordered map", node.start_mark,
344                    "expected a sequence, but found %s" % node.id, node.start_mark)
345        for subnode in node.value:
346            if not isinstance(subnode, MappingNode):
347                raise ConstructorError("while constructing an ordered map", node.start_mark,
348                        "expected a mapping of length 1, but found %s" % subnode.id,
349                        subnode.start_mark)
350            if len(subnode.value) != 1:
351                raise ConstructorError("while constructing an ordered map", node.start_mark,
352                        "expected a single mapping item, but found %d items" % len(subnode.value),
353                        subnode.start_mark)
354            key_node, value_node = subnode.value[0]
355            key = self.construct_object(key_node)
356            value = self.construct_object(value_node)
357            omap.append((key, value))
358
359    def construct_yaml_pairs(self, node):
360        # Note: the same code as `construct_yaml_omap`.
361        pairs = []
362        yield pairs
363        if not isinstance(node, SequenceNode):
364            raise ConstructorError("while constructing pairs", node.start_mark,
365                    "expected a sequence, but found %s" % node.id, node.start_mark)
366        for subnode in node.value:
367            if not isinstance(subnode, MappingNode):
368                raise ConstructorError("while constructing pairs", node.start_mark,
369                        "expected a mapping of length 1, but found %s" % subnode.id,
370                        subnode.start_mark)
371            if len(subnode.value) != 1:
372                raise ConstructorError("while constructing pairs", node.start_mark,
373                        "expected a single mapping item, but found %d items" % len(subnode.value),
374                        subnode.start_mark)
375            key_node, value_node = subnode.value[0]
376            key = self.construct_object(key_node)
377            value = self.construct_object(value_node)
378            pairs.append((key, value))
379
380    def construct_yaml_set(self, node):
381        data = set()
382        yield data
383        value = self.construct_mapping(node)
384        data.update(value)
385
386    def construct_yaml_str(self, node):
387        value = self.construct_scalar(node)
388        try:
389            return value.encode('ascii')
390        except UnicodeEncodeError:
391            return value
392
393    def construct_yaml_seq(self, node):
394        data = []
395        yield data
396        data.extend(self.construct_sequence(node))
397
398    def construct_yaml_map(self, node):
399        data = {}
400        yield data
401        value = self.construct_mapping(node)
402        data.update(value)
403
404    def construct_yaml_object(self, node, cls):
405        data = cls.__new__(cls)
406        yield data
407        if hasattr(data, '__setstate__'):
408            state = self.construct_mapping(node, deep=True)
409            data.__setstate__(state)
410        else:
411            state = self.construct_mapping(node)
412            data.__dict__.update(state)
413
414    def construct_undefined(self, node):
415        raise ConstructorError(None, None,
416                "could not determine a constructor for the tag %r" % node.tag.encode('utf-8'),
417                node.start_mark)
418
419SafeConstructor.add_constructor(
420        u'tag:yaml.org,2002:null',
421        SafeConstructor.construct_yaml_null)
422
423SafeConstructor.add_constructor(
424        u'tag:yaml.org,2002:bool',
425        SafeConstructor.construct_yaml_bool)
426
427SafeConstructor.add_constructor(
428        u'tag:yaml.org,2002:int',
429        SafeConstructor.construct_yaml_int)
430
431SafeConstructor.add_constructor(
432        u'tag:yaml.org,2002:float',
433        SafeConstructor.construct_yaml_float)
434
435SafeConstructor.add_constructor(
436        u'tag:yaml.org,2002:binary',
437        SafeConstructor.construct_yaml_binary)
438
439SafeConstructor.add_constructor(
440        u'tag:yaml.org,2002:timestamp',
441        SafeConstructor.construct_yaml_timestamp)
442
443SafeConstructor.add_constructor(
444        u'tag:yaml.org,2002:omap',
445        SafeConstructor.construct_yaml_omap)
446
447SafeConstructor.add_constructor(
448        u'tag:yaml.org,2002:pairs',
449        SafeConstructor.construct_yaml_pairs)
450
451SafeConstructor.add_constructor(
452        u'tag:yaml.org,2002:set',
453        SafeConstructor.construct_yaml_set)
454
455SafeConstructor.add_constructor(
456        u'tag:yaml.org,2002:str',
457        SafeConstructor.construct_yaml_str)
458
459SafeConstructor.add_constructor(
460        u'tag:yaml.org,2002:seq',
461        SafeConstructor.construct_yaml_seq)
462
463SafeConstructor.add_constructor(
464        u'tag:yaml.org,2002:map',
465        SafeConstructor.construct_yaml_map)
466
467SafeConstructor.add_constructor(None,
468        SafeConstructor.construct_undefined)
469
470class Constructor(SafeConstructor):
471
472    def construct_python_str(self, node):
473        return self.construct_scalar(node).encode('utf-8')
474
475    def construct_python_unicode(self, node):
476        return self.construct_scalar(node)
477
478    def construct_python_long(self, node):
479        return long(self.construct_yaml_int(node))
480
481    def construct_python_complex(self, node):
482       return complex(self.construct_scalar(node))
483
484    def construct_python_tuple(self, node):
485        return tuple(self.construct_sequence(node))
486
487    def find_python_module(self, name, mark):
488        if not name:
489            raise ConstructorError("while constructing a Python module", mark,
490                    "expected non-empty name appended to the tag", mark)
491        try:
492            __import__(name)
493        except ImportError, exc:
494            raise ConstructorError("while constructing a Python module", mark,
495                    "cannot find module %r (%s)" % (name.encode('utf-8'), exc), mark)
496        return sys.modules[name]
497
498    def find_python_name(self, name, mark):
499        if not name:
500            raise ConstructorError("while constructing a Python object", mark,
501                    "expected non-empty name appended to the tag", mark)
502        if u'.' in name:
503            # Python 2.4 only
504            #module_name, object_name = name.rsplit('.', 1)
505            items = name.split('.')
506            object_name = items.pop()
507            module_name = '.'.join(items)
508        else:
509            module_name = '__builtin__'
510            object_name = name
511        try:
512            __import__(module_name)
513        except ImportError, exc:
514            raise ConstructorError("while constructing a Python object", mark,
515                    "cannot find module %r (%s)" % (module_name.encode('utf-8'), exc), mark)
516        module = sys.modules[module_name]
517        if not hasattr(module, object_name):
518            raise ConstructorError("while constructing a Python object", mark,
519                    "cannot find %r in the module %r" % (object_name.encode('utf-8'),
520                        module.__name__), mark)
521        return getattr(module, object_name)
522
523    def construct_python_name(self, suffix, node):
524        value = self.construct_scalar(node)
525        if value:
526            raise ConstructorError("while constructing a Python name", node.start_mark,
527                    "expected the empty value, but found %r" % value.encode('utf-8'),
528                    node.start_mark)
529        return self.find_python_name(suffix, node.start_mark)
530
531    def construct_python_module(self, suffix, node):
532        value = self.construct_scalar(node)
533        if value:
534            raise ConstructorError("while constructing a Python module", node.start_mark,
535                    "expected the empty value, but found %r" % value.encode('utf-8'),
536                    node.start_mark)
537        return self.find_python_module(suffix, node.start_mark)
538
539    class classobj: pass
540
541    def make_python_instance(self, suffix, node,
542            args=None, kwds=None, newobj=False):
543        if not args:
544            args = []
545        if not kwds:
546            kwds = {}
547        cls = self.find_python_name(suffix, node.start_mark)
548        if newobj and isinstance(cls, type(self.classobj))  \
549                and not args and not kwds:
550            instance = self.classobj()
551            instance.__class__ = cls
552            return instance
553        elif newobj and isinstance(cls, type):
554            return cls.__new__(cls, *args, **kwds)
555        else:
556            return cls(*args, **kwds)
557
558    def set_python_instance_state(self, instance, state):
559        if hasattr(instance, '__setstate__'):
560            instance.__setstate__(state)
561        else:
562            slotstate = {}
563            if isinstance(state, tuple) and len(state) == 2:
564                state, slotstate = state
565            if hasattr(instance, '__dict__'):
566                instance.__dict__.update(state)
567            elif state:
568                slotstate.update(state)
569            for key, value in slotstate.items():
570                setattr(object, key, value)
571
572    def construct_python_object(self, suffix, node):
573        # Format:
574        #   !!python/object:module.name { ... state ... }
575        instance = self.make_python_instance(suffix, node, newobj=True)
576        yield instance
577        deep = hasattr(instance, '__setstate__')
578        state = self.construct_mapping(node, deep=deep)
579        self.set_python_instance_state(instance, state)
580
581    def construct_python_object_apply(self, suffix, node, newobj=False):
582        # Format:
583        #   !!python/object/apply       # (or !!python/object/new)
584        #   args: [ ... arguments ... ]
585        #   kwds: { ... keywords ... }
586        #   state: ... state ...
587        #   listitems: [ ... listitems ... ]
588        #   dictitems: { ... dictitems ... }
589        # or short format:
590        #   !!python/object/apply [ ... arguments ... ]
591        # The difference between !!python/object/apply and !!python/object/new
592        # is how an object is created, check make_python_instance for details.
593        if isinstance(node, SequenceNode):
594            args = self.construct_sequence(node, deep=True)
595            kwds = {}
596            state = {}
597            listitems = []
598            dictitems = {}
599        else:
600            value = self.construct_mapping(node, deep=True)
601            args = value.get('args', [])
602            kwds = value.get('kwds', {})
603            state = value.get('state', {})
604            listitems = value.get('listitems', [])
605            dictitems = value.get('dictitems', {})
606        instance = self.make_python_instance(suffix, node, args, kwds, newobj)
607        if state:
608            self.set_python_instance_state(instance, state)
609        if listitems:
610            instance.extend(listitems)
611        if dictitems:
612            for key in dictitems:
613                instance[key] = dictitems[key]
614        return instance
615
616    def construct_python_object_new(self, suffix, node):
617        return self.construct_python_object_apply(suffix, node, newobj=True)
618
619Constructor.add_constructor(
620    u'tag:yaml.org,2002:python/none',
621    Constructor.construct_yaml_null)
622
623Constructor.add_constructor(
624    u'tag:yaml.org,2002:python/bool',
625    Constructor.construct_yaml_bool)
626
627Constructor.add_constructor(
628    u'tag:yaml.org,2002:python/str',
629    Constructor.construct_python_str)
630
631Constructor.add_constructor(
632    u'tag:yaml.org,2002:python/unicode',
633    Constructor.construct_python_unicode)
634
635Constructor.add_constructor(
636    u'tag:yaml.org,2002:python/int',
637    Constructor.construct_yaml_int)
638
639Constructor.add_constructor(
640    u'tag:yaml.org,2002:python/long',
641    Constructor.construct_python_long)
642
643Constructor.add_constructor(
644    u'tag:yaml.org,2002:python/float',
645    Constructor.construct_yaml_float)
646
647Constructor.add_constructor(
648    u'tag:yaml.org,2002:python/complex',
649    Constructor.construct_python_complex)
650
651Constructor.add_constructor(
652    u'tag:yaml.org,2002:python/list',
653    Constructor.construct_yaml_seq)
654
655Constructor.add_constructor(
656    u'tag:yaml.org,2002:python/tuple',
657    Constructor.construct_python_tuple)
658
659Constructor.add_constructor(
660    u'tag:yaml.org,2002:python/dict',
661    Constructor.construct_yaml_map)
662
663Constructor.add_multi_constructor(
664    u'tag:yaml.org,2002:python/name:',
665    Constructor.construct_python_name)
666
667Constructor.add_multi_constructor(
668    u'tag:yaml.org,2002:python/module:',
669    Constructor.construct_python_module)
670
671Constructor.add_multi_constructor(
672    u'tag:yaml.org,2002:python/object:',
673    Constructor.construct_python_object)
674
675Constructor.add_multi_constructor(
676    u'tag:yaml.org,2002:python/object/apply:',
677    Constructor.construct_python_object_apply)
678
679Constructor.add_multi_constructor(
680    u'tag:yaml.org,2002:python/object/new:',
681    Constructor.construct_python_object_new)
682
Note: See TracBrowser for help on using the repository browser.