copyreg.py 6.7 KB

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