Ignore:
Timestamp:
04/22/06 16:40:43 (8 years ago)
Author:
xi
Message:

Add support for pickling/unpickling python objects.

File:
1 edited

Legend:

Unmodified
Added
Removed
  • pyyaml/trunk/lib/yaml/representer.py

    r143 r147  
    1717    from sets import Set as set 
    1818 
    19 import sys 
     19import sys, copy_reg 
    2020 
    2121class RepresenterError(YAMLError): 
     
    2525 
    2626    yaml_representers = {} 
     27    yaml_multi_representers = {} 
    2728 
    2829    def __init__(self): 
     
    3031 
    3132    def represent(self, data): 
    32         node = self.represent_object(data) 
     33        node = self.represent_data(data) 
    3334        self.serialize(node) 
    3435        self.represented_objects = {} 
     
    5051        return bases 
    5152 
    52     def represent_object(self, data): 
     53    def represent_data(self, data): 
    5354        if self.ignore_aliases(data): 
    5455            alias_key = None 
     
    6566        if type(data) is self.instance_type: 
    6667            data_types = self.get_classobj_bases(data.__class__)+list(data_types) 
    67         for data_type in data_types: 
    68             if data_type in self.yaml_representers: 
    69                 node = self.yaml_representers[data_type](self, data) 
    70                 break 
    71         else: 
    72             if None in self.yaml_representers: 
    73                 node = self.yaml_representers[None](self, data) 
     68        if data_types[0] in self.yaml_representers: 
     69            node = self.yaml_representers[data_types[0]](self, data) 
     70        else: 
     71            for data_type in data_types: 
     72                if data_type in self.yaml_multi_representers: 
     73                    node = self.yaml_multi_representers[data_type](self, data) 
     74                    break 
    7475            else: 
    75                 node = ScalarNode(None, unicode(data)) 
     76                if None in self.yaml_multi_representers: 
     77                    node = self.yaml_multi_representers[None](self, data) 
     78                elif None in self.yaml_representers: 
     79                    node = self.yaml_representers[None](self, data) 
     80                else: 
     81                    node = ScalarNode(None, unicode(data)) 
    7682        if alias_key is not None: 
    7783            self.represented_objects[alias_key] = node 
     
    8490    add_representer = classmethod(add_representer) 
    8591 
     92    def add_multi_representer(cls, data_type, representer): 
     93        if not 'yaml_multi_representers' in cls.__dict__: 
     94            cls.yaml_multi_representers = cls.yaml_multi_representers.copy() 
     95        cls.yaml_multi_representers[data_type] = representer 
     96    add_multi_representer = classmethod(add_multi_representer) 
     97 
    8698    def represent_scalar(self, tag, value, style=None): 
    8799        return ScalarNode(tag, value, style=style) 
    88100 
    89101    def represent_sequence(self, tag, sequence, flow_style=None): 
     102        best_style = True 
    90103        value = [] 
    91104        for item in sequence: 
    92             value.append(self.represent_object(item)) 
     105            node_item = self.represent_data(item) 
     106            if not (isinstance(node_item, ScalarNode) and not node_item.style): 
     107                best_style = False 
     108            value.append(self.represent_data(item)) 
     109        if flow_style is None: 
     110            flow_style = best_style 
    93111        return SequenceNode(tag, value, flow_style=flow_style) 
    94112 
    95113    def represent_mapping(self, tag, mapping, flow_style=None): 
     114        best_style = True 
    96115        if hasattr(mapping, 'keys'): 
    97116            value = {} 
    98117            for item_key in mapping.keys(): 
    99118                item_value = mapping[item_key] 
    100                 value[self.represent_object(item_key)] =    \ 
    101                         self.represent_object(item_value) 
     119                node_key = self.represent_data(item_key) 
     120                node_value = self.represent_data(item_value) 
     121                if not (isinstance(node_key, ScalarNode) and not node_key.style): 
     122                    best_style = False 
     123                if not (isinstance(node_value, ScalarNode) and not node_value.style): 
     124                    best_style = False 
     125                value[node_key] = node_value 
    102126        else: 
    103127            value = [] 
    104128            for item_key, item_value in mapping: 
    105                 value.append((self.represent_object(item_key), 
    106                         self.represent_object(item_value))) 
     129                node_key = self.represent_data(item_key) 
     130                node_value = self.represent_data(item_value) 
     131                if not (isinstance(node_key, ScalarNode) and not node_key.style): 
     132                    best_style = False 
     133                if not (isinstance(node_value, ScalarNode) and not node_value.style): 
     134                    best_style = False 
     135                value.append((node_key, node_value)) 
     136        if flow_style is None: 
     137            flow_style = best_style 
    107138        return MappingNode(tag, value, flow_style=flow_style) 
    108139 
     
    259290 
    260291class Representer(SafeRepresenter): 
    261      
     292 
    262293    def represent_str(self, data): 
    263294        tag = None 
     
    313344                u'tag:yaml.org,2002:python/module:'+data.__name__, u'') 
    314345 
     346    def represent_instance(self, data): 
     347        # For instances of classic classes, we use __getinitargs__ and 
     348        # __getstate__ to serialize the data. 
     349 
     350        # If data.__getinitargs__ exists, the object must be reconstructed by 
     351        # calling cls(**args), where args is a tuple returned by 
     352        # __getinitargs__. Otherwise, the cls.__init__ method should never be 
     353        # called and the class instance is created by instantiating a trivial 
     354        # class and assigning to the instance's __class__ variable. 
     355 
     356        # If data.__getstate__ exists, it returns the state of the object. 
     357        # Otherwise, the state of the object is data.__dict__. 
     358 
     359        # We produce either a !!python/object or !!python/object/new node. 
     360        # If data.__getinitargs__ does not exist and state is a dictionary, we 
     361        # produce a !!python/object node . Otherwise we produce a 
     362        # !!python/object/new node. 
     363 
     364        cls = data.__class__ 
     365        class_name = u'%s.%s' % (cls.__module__, cls.__name__) 
     366        args = None 
     367        state = None 
     368        if hasattr(data, '__getinitargs__'): 
     369            args = list(data.__getinitargs__()) 
     370        if hasattr(data, '__getstate__'): 
     371            state = data.__getstate__() 
     372        else: 
     373            state = data.__dict__ 
     374        if args is None and isinstance(state, dict): 
     375            return self.represent_mapping( 
     376                    u'tag:yaml.org,2002:python/object:'+class_name, state) 
     377        if isinstance(state, dict) and not state: 
     378            return self.represent_sequence( 
     379                    u'tag:yaml.org,2002:python/object/new:'+class_name, args) 
     380        value = {} 
     381        if args: 
     382            value['args'] = args 
     383        value['state'] = state 
     384        return self.represent_mapping( 
     385                u'tag:yaml.org,2002:python/object/new:'+class_name, value) 
     386 
     387    def represent_object(self, data): 
     388        # We use __reduce__ API to save the data. data.__reduce__ returns 
     389        # a tuple of length 2-5: 
     390        #   (function, args, state, listitems, dictitems) 
     391 
     392        # For reconstructing, we calls function(*args), then set its state, 
     393        # listitems, and dictitems if they are not None. 
     394 
     395        # A special case is when function.__name__ == '__newobj__'. In this 
     396        # case we create the object with args[0].__new__(*args). 
     397 
     398        # Another special case is when __reduce__ returns a string - we don't 
     399        # support it. 
     400 
     401        # We produce a !!python/object, !!python/object/new or 
     402        # !!python/object/apply node. 
     403 
     404        cls = type(data) 
     405        if cls in copy_reg.dispatch_table: 
     406            reduce = copy_reg.dispatch_table[cls] 
     407        elif hasattr(data, '__reduce_ex__'): 
     408            reduce = data.__reduce_ex__(2) 
     409        elif hasattr(data, '__reduce__'): 
     410            reduce = data.__reduce__() 
     411        else: 
     412            raise RepresenterError("cannot represent object: %r" % data) 
     413        reduce = (list(reduce)+[None]*5)[:5] 
     414        function, args, state, listitems, dictitems = reduce 
     415        args = list(args) 
     416        if state is None: 
     417            state = {} 
     418        if listitems is not None: 
     419            listitems = list(listitems) 
     420        if dictitems is not None: 
     421            dictitems = dict(dictitems) 
     422        if function.__name__ == '__newobj__': 
     423            function = args[0] 
     424            args = args[1:] 
     425            tag = u'tag:yaml.org,2002:python/object/new:' 
     426            newobj = True 
     427        else: 
     428            tag = u'tag:yaml.org,2002:python/object/apply:' 
     429            newobj = False 
     430        function_name = u'%s.%s' % (function.__module__, function.__name__) 
     431        if not args and not listitems and not dictitems \ 
     432                and isinstance(state, dict) and newobj: 
     433            return self.represent_mapping( 
     434                    u'tag:yaml.org,2002:python/object:'+function_name, state) 
     435        if not listitems and not dictitems  \ 
     436                and isinstance(state, dict) and not state: 
     437            return self.represent_sequence(tag+function_name, args) 
     438        value = {} 
     439        if args: 
     440            value['args'] = args 
     441        if state or not isinstance(state, dict): 
     442            value['state'] = state 
     443        if listitems: 
     444            value['listitems'] = listitems 
     445        if dictitems: 
     446            value['dictitems'] = dictitems 
     447        return self.represent_mapping(tag+function_name, value) 
     448 
    315449Representer.add_representer(str, 
    316450        Representer.represent_str) 
     
    343477        Representer.represent_module) 
    344478 
     479Representer.add_multi_representer(Representer.instance_type, 
     480        Representer.represent_instance) 
     481 
     482Representer.add_multi_representer(object, 
     483        Representer.represent_object) 
     484 
Note: See TracChangeset for help on using the changeset viewer.