183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh"""Helper to provide extensibility for pickle/cPickle.
283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew HsiehThis is only useful to add pickle support for extension types defined in
483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew HsiehC, not for instances of user-defined classes.
583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh"""
683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom types import ClassType as _ClassType
883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh__all__ = ["pickle", "constructor",
1083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh           "add_extension", "remove_extension", "clear_extension_cache"]
1183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdispatch_table = {}
1383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef pickle(ob_type, pickle_function, constructor_ob=None):
1583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if type(ob_type) is _ClassType:
1683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise TypeError("copy_reg is not intended for use with classes")
1783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
1883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if not hasattr(pickle_function, '__call__'):
1983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise TypeError("reduction functions must be callable")
2083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    dispatch_table[ob_type] = pickle_function
2183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
2283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # The constructor_ob function is a vestige of safe for unpickling.
2383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # There is no reason for the caller to pass it anymore.
2483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if constructor_ob is not None:
2583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        constructor(constructor_ob)
2683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
2783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef constructor(object):
2883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if not hasattr(object, '__call__'):
2983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise TypeError("constructors must be callable")
3083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
3183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Example: provide pickling support for complex numbers.
3283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
3383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehtry:
3483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    complex
3583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehexcept NameError:
3683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    pass
3783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehelse:
3883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
3983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def pickle_complex(c):
4083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return complex, (c.real, c.imag)
4183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    pickle(complex, pickle_complex, complex)
4383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Support for pickling new-style objects
4583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _reconstructor(cls, base, state):
4783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if base is object:
4883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        obj = object.__new__(cls)
4983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    else:
5083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        obj = base.__new__(cls, state)
5183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if base.__init__ != object.__init__:
5283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            base.__init__(obj, state)
5383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return obj
5483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh_HEAPTYPE = 1<<9
5683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Python code for object.__reduce_ex__ for protocols 0 and 1
5883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _reduce_ex(self, proto):
6083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    assert proto < 2
6183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    for base in self.__class__.__mro__:
6283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if hasattr(base, '__flags__') and not base.__flags__ & _HEAPTYPE:
6383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            break
6483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    else:
6583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        base = object # not really reachable
6683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if base is object:
6783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        state = None
6883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    else:
6983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if base is self.__class__:
7083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            raise TypeError, "can't pickle %s objects" % base.__name__
7183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        state = base(self)
7283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    args = (self.__class__, base, state)
7383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    try:
7483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        getstate = self.__getstate__
7583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    except AttributeError:
7683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if getattr(self, "__slots__", None):
7783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            raise TypeError("a class that defines __slots__ without "
7883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            "defining __getstate__ cannot be pickled")
7983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        try:
8083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            dict = self.__dict__
8183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        except AttributeError:
8283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            dict = None
8383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    else:
8483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        dict = getstate()
8583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if dict:
8683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return _reconstructor, args, dict
8783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    else:
8883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return _reconstructor, args
8983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
9083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Helper for __reduce_ex__ protocol 2
9183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
9283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef __newobj__(cls, *args):
9383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return cls.__new__(cls, *args)
9483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
9583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _slotnames(cls):
9683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """Return a list of slot names for a given class.
9783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
9883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    This needs to find slots defined by the class and its bases, so we
9983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    can't simply return the __slots__ attribute.  We must walk down
10083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    the Method Resolution Order and concatenate the __slots__ of each
10183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    class found there.  (This assumes classes don't modify their
10283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    __slots__ attribute to misrepresent their slots after the class is
10383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    defined.)
10483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """
10583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
10683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # Get the value from a cache in the class if possible
10783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    names = cls.__dict__.get("__slotnames__")
10883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if names is not None:
10983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return names
11083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
11183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # Not cached -- calculate the value
11283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    names = []
11383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if not hasattr(cls, "__slots__"):
11483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # This class has no slots
11583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        pass
11683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    else:
11783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        # Slots found -- gather slot names from all base classes
11883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        for c in cls.__mro__:
11983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if "__slots__" in c.__dict__:
12083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                slots = c.__dict__['__slots__']
12183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                # if class has a single slot, it can be given as a string
12283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if isinstance(slots, basestring):
12383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    slots = (slots,)
12483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                for name in slots:
12583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    # special descriptors
12683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if name in ("__dict__", "__weakref__"):
12783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        continue
12883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    # mangled names
12983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    elif name.startswith('__') and not name.endswith('__'):
13083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        names.append('_%s%s' % (c.__name__, name))
13183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    else:
13283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        names.append(name)
13383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
13483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # Cache the outcome in the class if at all possible
13583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    try:
13683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        cls.__slotnames__ = names
13783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    except:
13883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        pass # But don't die if we can't
13983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
14083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return names
14183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
14283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# A registry of extension codes.  This is an ad-hoc compression
14383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# mechanism.  Whenever a global reference to <module>, <name> is about
14483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# to be pickled, the (<module>, <name>) tuple is looked up here to see
14583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# if it is a registered extension code for it.  Extension codes are
14683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# universal, so that the meaning of a pickle does not depend on
14783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# context.  (There are also some codes reserved for local use that
14883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# don't have this restriction.)  Codes are positive ints; 0 is
14983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# reserved.
15083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
15183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh_extension_registry = {}                # key -> code
15283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh_inverted_registry = {}                 # code -> key
15383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh_extension_cache = {}                   # code -> object
15483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Don't ever rebind those names:  cPickle grabs a reference to them when
15583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# it's initialized, and won't see a rebinding.
15683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
15783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef add_extension(module, name, code):
15883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """Register an extension code."""
15983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    code = int(code)
16083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if not 1 <= code <= 0x7fffffff:
16183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise ValueError, "code out of range"
16283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    key = (module, name)
16383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if (_extension_registry.get(key) == code and
16483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        _inverted_registry.get(code) == key):
16583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return # Redundant registrations are benign
16683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if key in _extension_registry:
16783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise ValueError("key %s is already registered with code %s" %
16883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                         (key, _extension_registry[key]))
16983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if code in _inverted_registry:
17083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise ValueError("code %s is already in use for key %s" %
17183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                         (code, _inverted_registry[code]))
17283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    _extension_registry[key] = code
17383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    _inverted_registry[code] = key
17483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef remove_extension(module, name, code):
17683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    """Unregister an extension code.  For testing only."""
17783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    key = (module, name)
17883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if (_extension_registry.get(key) != code or
17983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        _inverted_registry.get(code) != key):
18083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise ValueError("key %s is not registered with code %s" %
18183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                         (key, code))
18283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    del _extension_registry[key]
18383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    del _inverted_registry[code]
18483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if code in _extension_cache:
18583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        del _extension_cache[code]
18683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
18783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef clear_extension_cache():
18883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    _extension_cache.clear()
18983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
19083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Standard extension code assignments
19183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
19283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Reserved ranges
19383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
19483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# First  Last Count  Purpose
19583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#     1   127   127  Reserved for Python standard library
19683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#   128   191    64  Reserved for Zope
19783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#   192   239    48  Reserved for 3rd parties
19883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#   240   255    16  Reserved for private use (will never be assigned)
19983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#   256   Inf   Inf  Reserved for future assignment
20083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
20183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Extension codes are assigned by the Python Software Foundation.
202