Lib/copy_reg.py

   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.