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