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