Home | History | Annotate | Download | only in python2.7 
      1 """Helper to provide extensibility for pickle/cPickle.
      2 
      3 This is only useful to add pickle support for extension types defined in
      4 C, not for instances of user-defined classes.
      5 """
      6 
      7 from types import ClassType as _ClassType
      8 
      9 __all__ = ["pickle", "constructor",
     10            "add_extension", "remove_extension", "clear_extension_cache"]
     11 
     12 dispatch_table = {}
     13 
     14 def pickle(ob_type, pickle_function, constructor_ob=None):
     15     if type(ob_type) is _ClassType:
     16         raise TypeError("copy_reg is not intended for use with classes")
     17 
     18     if not hasattr(pickle_function, '__call__'):
     19         raise TypeError("reduction functions must be callable")
     20     dispatch_table[ob_type] = pickle_function
     21 
     22     # The constructor_ob function is a vestige of safe for unpickling.
     23     # There is no reason for the caller to pass it anymore.
     24     if constructor_ob is not None:
     25         constructor(constructor_ob)
     26 
     27 def constructor(object):
     28     if not hasattr(object, '__call__'):
     29         raise TypeError("constructors must be callable")
     30 
     31 # Example: provide pickling support for complex numbers.
     32 
     33 try:
     34     complex
     35 except NameError:
     36     pass
     37 else:
     38 
     39     def pickle_complex(c):
     40         return complex, (c.real, c.imag)
     41 
     42     pickle(complex, pickle_complex, complex)
     43 
     44 # Support for pickling new-style objects
     45 
     46 def _reconstructor(cls, base, state):
     47     if base is object:
     48         obj = object.__new__(cls)
     49     else:
     50         obj = base.__new__(cls, state)
     51         if base.__init__ != object.__init__:
     52             base.__init__(obj, state)
     53     return obj
     54 
     55 _HEAPTYPE = 1<<9
     56 
     57 # Python code for object.__reduce_ex__ for protocols 0 and 1
     58 
     59 def _reduce_ex(self, proto):
     60     assert proto < 2
     61     for base in self.__class__.__mro__:
     62         if hasattr(base, '__flags__') and not base.__flags__ & _HEAPTYPE:
     63             break
     64     else:
     65         base = object # not really reachable
     66     if base is object:
     67         state = None
     68     else:
     69         if base is self.__class__:
     70             raise TypeError, "can't pickle %s objects" % base.__name__
     71         state = base(self)
     72     args = (self.__class__, base, state)
     73     try:
     74         getstate = self.__getstate__
     75     except AttributeError:
     76         if getattr(self, "__slots__", None):
     77             raise TypeError("a class that defines __slots__ without "
     78                             "defining __getstate__ cannot be pickled")
     79         try:
     80             dict = self.__dict__
     81         except AttributeError:
     82             dict = None
     83     else:
     84         dict = getstate()
     85     if dict:
     86         return _reconstructor, args, dict
     87     else:
     88         return _reconstructor, args
     89 
     90 # Helper for __reduce_ex__ protocol 2
     91 
     92 def __newobj__(cls, *args):
     93     return cls.__new__(cls, *args)
     94 
     95 def _slotnames(cls):
     96     """Return a list of slot names for a given class.
     97 
     98     This needs to find slots defined by the class and its bases, so we
     99     can't simply return the __slots__ attribute.  We must walk down
    100     the Method Resolution Order and concatenate the __slots__ of each
    101     class found there.  (This assumes classes don't modify their
    102     __slots__ attribute to misrepresent their slots after the class is
    103     defined.)
    104     """
    105 
    106     # Get the value from a cache in the class if possible
    107     names = cls.__dict__.get("__slotnames__")
    108     if names is not None:
    109         return names
    110 
    111     # Not cached -- calculate the value
    112     names = []
    113     if not hasattr(cls, "__slots__"):
    114         # This class has no slots
    115         pass
    116     else:
    117         # Slots found -- gather slot names from all base classes
    118         for c in cls.__mro__:
    119             if "__slots__" in c.__dict__:
    120                 slots = c.__dict__['__slots__']
    121                 # if class has a single slot, it can be given as a string
    122                 if isinstance(slots, basestring):
    123                     slots = (slots,)
    124                 for name in slots:
    125                     # special descriptors
    126                     if name in ("__dict__", "__weakref__"):
    127                         continue
    128                     # mangled names
    129                     elif name.startswith('__') and not name.endswith('__'):
    130                         names.append('_%s%s' % (c.__name__, name))
    131                     else:
    132                         names.append(name)
    133 
    134     # Cache the outcome in the class if at all possible
    135     try:
    136         cls.__slotnames__ = names
    137     except:
    138         pass # But don't die if we can't
    139 
    140     return names
    141 
    142 # A registry of extension codes.  This is an ad-hoc compression
    143 # mechanism.  Whenever a global reference to <module>, <name> is about
    144 # to be pickled, the (<module>, <name>) tuple is looked up here to see
    145 # if it is a registered extension code for it.  Extension codes are
    146 # universal, so that the meaning of a pickle does not depend on
    147 # context.  (There are also some codes reserved for local use that
    148 # don't have this restriction.)  Codes are positive ints; 0 is
    149 # reserved.
    150 
    151 _extension_registry = {}                # key -> code
    152 _inverted_registry = {}                 # code -> key
    153 _extension_cache = {}                   # code -> object
    154 # Don't ever rebind those names:  cPickle grabs a reference to them when
    155 # it's initialized, and won't see a rebinding.
    156 
    157 def add_extension(module, name, code):
    158     """Register an extension code."""
    159     code = int(code)
    160     if not 1 <= code <= 0x7fffffff:
    161         raise ValueError, "code out of range"
    162     key = (module, name)
    163     if (_extension_registry.get(key) == code and
    164         _inverted_registry.get(code) == key):
    165         return # Redundant registrations are benign
    166     if key in _extension_registry:
    167         raise ValueError("key %s is already registered with code %s" %
    168                          (key, _extension_registry[key]))
    169     if code in _inverted_registry:
    170         raise ValueError("code %s is already in use for key %s" %
    171                          (code, _inverted_registry[code]))
    172     _extension_registry[key] = code
    173     _inverted_registry[code] = key
    174 
    175 def remove_extension(module, name, code):
    176     """Unregister an extension code.  For testing only."""
    177     key = (module, name)
    178     if (_extension_registry.get(key) != code or
    179         _inverted_registry.get(code) != key):
    180         raise ValueError("key %s is not registered with code %s" %
    181                          (key, code))
    182     del _extension_registry[key]
    183     del _inverted_registry[code]
    184     if code in _extension_cache:
    185         del _extension_cache[code]
    186 
    187 def clear_extension_cache():
    188     _extension_cache.clear()
    189 
    190 # Standard extension code assignments
    191 
    192 # Reserved ranges
    193 
    194 # First  Last Count  Purpose
    195 #     1   127   127  Reserved for Python standard library
    196 #   128   191    64  Reserved for Zope
    197 #   192   239    48  Reserved for 3rd parties
    198 #   240   255    16  Reserved for private use (will never be assigned)
    199 #   256   Inf   Inf  Reserved for future assignment
    200 
    201 # Extension codes are assigned by the Python Software Foundation.
    202