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
7from types import ClassType as _ClassType
8
9__all__ = ["pickle", "constructor",
10           "add_extension", "remove_extension", "clear_extension_cache"]
11
12dispatch_table = {}
13
14def 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
27def 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
33try:
34    complex
35except NameError:
36    pass
37else:
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
46def _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
59def _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
92def __newobj__(cls, *args):
93    return cls.__new__(cls, *args)
94
95def _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                        stripped = c.__name__.lstrip('_')
131                        if stripped:
132                            names.append('_%s%s' % (stripped, name))
133                        else:
134                            names.append(name)
135                    else:
136                        names.append(name)
137
138    # Cache the outcome in the class if at all possible
139    try:
140        cls.__slotnames__ = names
141    except:
142        pass # But don't die if we can't
143
144    return names
145
146# A registry of extension codes.  This is an ad-hoc compression
147# mechanism.  Whenever a global reference to <module>, <name> is about
148# to be pickled, the (<module>, <name>) tuple is looked up here to see
149# if it is a registered extension code for it.  Extension codes are
150# universal, so that the meaning of a pickle does not depend on
151# context.  (There are also some codes reserved for local use that
152# don't have this restriction.)  Codes are positive ints; 0 is
153# reserved.
154
155_extension_registry = {}                # key -> code
156_inverted_registry = {}                 # code -> key
157_extension_cache = {}                   # code -> object
158# Don't ever rebind those names:  cPickle grabs a reference to them when
159# it's initialized, and won't see a rebinding.
160
161def add_extension(module, name, code):
162    """Register an extension code."""
163    code = int(code)
164    if not 1 <= code <= 0x7fffffff:
165        raise ValueError, "code out of range"
166    key = (module, name)
167    if (_extension_registry.get(key) == code and
168        _inverted_registry.get(code) == key):
169        return # Redundant registrations are benign
170    if key in _extension_registry:
171        raise ValueError("key %s is already registered with code %s" %
172                         (key, _extension_registry[key]))
173    if code in _inverted_registry:
174        raise ValueError("code %s is already in use for key %s" %
175                         (code, _inverted_registry[code]))
176    _extension_registry[key] = code
177    _inverted_registry[code] = key
178
179def remove_extension(module, name, code):
180    """Unregister an extension code.  For testing only."""
181    key = (module, name)
182    if (_extension_registry.get(key) != code or
183        _inverted_registry.get(code) != key):
184        raise ValueError("key %s is not registered with code %s" %
185                         (key, code))
186    del _extension_registry[key]
187    del _inverted_registry[code]
188    if code in _extension_cache:
189        del _extension_cache[code]
190
191def clear_extension_cache():
192    _extension_cache.clear()
193
194# Standard extension code assignments
195
196# Reserved ranges
197
198# First  Last Count  Purpose
199#     1   127   127  Reserved for Python standard library
200#   128   191    64  Reserved for Zope
201#   192   239    48  Reserved for 3rd parties
202#   240   255    16  Reserved for private use (will never be assigned)
203#   256   Inf   Inf  Reserved for future assignment
204
205# Extension codes are assigned by the Python Software Foundation.
206