/Lib/copy_reg.py

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