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