14710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm"""Helper to provide extensibility for pickle/cPickle.
24710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
34710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmThis is only useful to add pickle support for extension types defined in
44710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmC, not for instances of user-defined classes.
54710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm"""
64710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
74710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom types import ClassType as _ClassType
84710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
94710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm__all__ = ["pickle", "constructor",
104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm           "add_extension", "remove_extension", "clear_extension_cache"]
114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdispatch_table = {}
134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef pickle(ob_type, pickle_function, constructor_ob=None):
154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if type(ob_type) is _ClassType:
164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        raise TypeError("copy_reg is not intended for use with classes")
174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if not hasattr(pickle_function, '__call__'):
194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        raise TypeError("reduction functions must be callable")
204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    dispatch_table[ob_type] = pickle_function
214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    # The constructor_ob function is a vestige of safe for unpickling.
234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    # There is no reason for the caller to pass it anymore.
244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if constructor_ob is not None:
254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        constructor(constructor_ob)
264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef constructor(object):
284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if not hasattr(object, '__call__'):
294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        raise TypeError("constructors must be callable")
304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Example: provide pickling support for complex numbers.
324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmtry:
344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    complex
354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmexcept NameError:
364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    pass
374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmelse:
384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    def pickle_complex(c):
404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        return complex, (c.real, c.imag)
414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    pickle(complex, pickle_complex, complex)
434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Support for pickling new-style objects
454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef _reconstructor(cls, base, state):
474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if base is object:
484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        obj = object.__new__(cls)
494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    else:
504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        obj = base.__new__(cls, state)
514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        if base.__init__ != object.__init__:
524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm            base.__init__(obj, state)
534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    return obj
544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm_HEAPTYPE = 1<<9
564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Python code for object.__reduce_ex__ for protocols 0 and 1
584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef _reduce_ex(self, proto):
604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    assert proto < 2
614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    for base in self.__class__.__mro__:
624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        if hasattr(base, '__flags__') and not base.__flags__ & _HEAPTYPE:
634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm            break
644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    else:
654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        base = object # not really reachable
664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if base is object:
674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        state = None
684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    else:
694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        if base is self.__class__:
704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm            raise TypeError, "can't pickle %s objects" % base.__name__
714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        state = base(self)
724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    args = (self.__class__, base, state)
734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    try:
744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        getstate = self.__getstate__
754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    except AttributeError:
764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        if getattr(self, "__slots__", None):
774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm            raise TypeError("a class that defines __slots__ without "
784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                            "defining __getstate__ cannot be pickled")
794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        try:
804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm            dict = self.__dict__
814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        except AttributeError:
824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm            dict = None
834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    else:
844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        dict = getstate()
854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if dict:
864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        return _reconstructor, args, dict
874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    else:
884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        return _reconstructor, args
894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Helper for __reduce_ex__ protocol 2
914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef __newobj__(cls, *args):
934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    return cls.__new__(cls, *args)
944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef _slotnames(cls):
964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    """Return a list of slot names for a given class.
974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    This needs to find slots defined by the class and its bases, so we
994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    can't simply return the __slots__ attribute.  We must walk down
1004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    the Method Resolution Order and concatenate the __slots__ of each
1014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    class found there.  (This assumes classes don't modify their
1024710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    __slots__ attribute to misrepresent their slots after the class is
1034710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    defined.)
1044710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    """
1054710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1064710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    # Get the value from a cache in the class if possible
1074710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    names = cls.__dict__.get("__slotnames__")
1084710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if names is not None:
1094710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        return names
1104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    # Not cached -- calculate the value
1124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    names = []
1134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if not hasattr(cls, "__slots__"):
1144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        # This class has no slots
1154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        pass
1164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    else:
1174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        # Slots found -- gather slot names from all base classes
1184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        for c in cls.__mro__:
1194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm            if "__slots__" in c.__dict__:
1204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                slots = c.__dict__['__slots__']
1214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                # if class has a single slot, it can be given as a string
1224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                if isinstance(slots, basestring):
1234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                    slots = (slots,)
1244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                for name in slots:
1254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                    # special descriptors
1264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                    if name in ("__dict__", "__weakref__"):
1274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                        continue
1284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                    # mangled names
1294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                    elif name.startswith('__') and not name.endswith('__'):
1304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                        names.append('_%s%s' % (c.__name__, name))
1314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                    else:
1324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                        names.append(name)
1334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    # Cache the outcome in the class if at all possible
1354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    try:
1364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        cls.__slotnames__ = names
1374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    except:
1384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        pass # But don't die if we can't
1394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    return names
1414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# A registry of extension codes.  This is an ad-hoc compression
1434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# mechanism.  Whenever a global reference to <module>, <name> is about
1444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# to be pickled, the (<module>, <name>) tuple is looked up here to see
1454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# if it is a registered extension code for it.  Extension codes are
1464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# universal, so that the meaning of a pickle does not depend on
1474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# context.  (There are also some codes reserved for local use that
1484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# don't have this restriction.)  Codes are positive ints; 0 is
1494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# reserved.
1504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm_extension_registry = {}                # key -> code
1524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm_inverted_registry = {}                 # code -> key
1534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm_extension_cache = {}                   # code -> object
1544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Don't ever rebind those names:  cPickle grabs a reference to them when
1554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# it's initialized, and won't see a rebinding.
1564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef add_extension(module, name, code):
1584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    """Register an extension code."""
1594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    code = int(code)
1604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if not 1 <= code <= 0x7fffffff:
1614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        raise ValueError, "code out of range"
1624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    key = (module, name)
1634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if (_extension_registry.get(key) == code and
1644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        _inverted_registry.get(code) == key):
1654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        return # Redundant registrations are benign
1664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if key in _extension_registry:
1674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        raise ValueError("key %s is already registered with code %s" %
1684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                         (key, _extension_registry[key]))
1694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if code in _inverted_registry:
1704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        raise ValueError("code %s is already in use for key %s" %
1714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                         (code, _inverted_registry[code]))
1724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    _extension_registry[key] = code
1734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    _inverted_registry[code] = key
1744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef remove_extension(module, name, code):
1764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    """Unregister an extension code.  For testing only."""
1774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    key = (module, name)
1784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if (_extension_registry.get(key) != code or
1794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        _inverted_registry.get(code) != key):
1804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        raise ValueError("key %s is not registered with code %s" %
1814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm                         (key, code))
1824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    del _extension_registry[key]
1834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    del _inverted_registry[code]
1844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    if code in _extension_cache:
1854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm        del _extension_cache[code]
1864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef clear_extension_cache():
1884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm    _extension_cache.clear()
1894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Standard extension code assignments
1914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Reserved ranges
1934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
1944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# First  Last Count  Purpose
1954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm#     1   127   127  Reserved for Python standard library
1964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm#   128   191    64  Reserved for Zope
1974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm#   192   239    48  Reserved for 3rd parties
1984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm#   240   255    16  Reserved for private use (will never be assigned)
1994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm#   256   Inf   Inf  Reserved for future assignment
2004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm
2014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Extension codes are assigned by the Python Software Foundation.
202