copyreg.py 7.4 KB

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