copy_reg.py 6.8 KB

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