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