copy_reg.py 6.8 KB

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