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 def pickle_complex(c):
29 return complex, (c.real, c.imag)
30
31 pickle(complex, pickle_complex, complex)
32
33 def pickle_union(obj):
34 import functools, operator
35 return functools.reduce, (operator.or_, obj.__args__)
36
37 pickle(type(int | str), pickle_union)
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 _new_type = type(int.__new__)
52
53 # Python code for object.__reduce_ex__ for protocols 0 and 1
54
55 def _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
98 def __newobj__(cls, *args):
99 return cls.__new__(cls, *args)
100
101 def __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
107 def _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
173 def 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
191 def 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
203 def 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.