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