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

Revision 147, 23.4 KB checked in by xi, 9 years ago (diff)

Add support for pickling/unpickling python objects.

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