1 """Helper to provide extensibility for pickle/cPickle.

     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 from types import ClassType as _ClassType

     8 

     9 __all__ = ["pickle", "constructor",

    10            "add_extension", "remove_extension", "clear_extension_cache"]

    11 

    12 dispatch_table = {}

    13 

    14 def pickle(ob_type, pickle_function, constructor_ob=None):

    15     if type(ob_type) is _ClassType:

    16         raise TypeError("copy_reg is not intended for use with classes")

    17 

    18     if not hasattr(pickle_function, '__call__'):

    19         raise TypeError("reduction functions must be callable")

    20     dispatch_table[ob_type] = pickle_function

    21 

    22     # The constructor_ob function is a vestige of safe for unpickling.

    23     # There is no reason for the caller to pass it anymore.

    24     if constructor_ob is not None:

    25         constructor(constructor_ob)

    26 

    27 def constructor(object):

    28     if not hasattr(object, '__call__'):

    29         raise TypeError("constructors must be callable")

    30 

    31 # Example: provide pickling support for complex numbers.

    32 

    33 try:

    34     complex

    35 except NameError:

    36     pass

    37 else:

    38 

    39     def pickle_complex(c):

    40         return complex, (c.real, c.imag)

    41 

    42     pickle(complex, pickle_complex, complex)

    43 

    44 # Support for pickling new-style objects

    45 

    46 def _reconstructor(cls, base, state):

    47     if base is object:

    48         obj = object.__new__(cls)

    49     else:

    50         obj = base.__new__(cls, state)

    51         if base.__init__ != object.__init__:

    52             base.__init__(obj, state)

    53     return obj

    54 

    55 _HEAPTYPE = 1<<9

    56 

    57 # Python code for object.__reduce_ex__ for protocols 0 and 1

    58 

    59 def _reduce_ex(self, proto):

    60     assert proto < 2

    61     for base in self.__class__.__mro__:

    62         if hasattr(base, '__flags__') and not base.__flags__ & _HEAPTYPE:

    63             break

    64     else:

    65         base = object # not really reachable

    66     if base is object:

    67         state = None

    68     else:

    69         if base is self.__class__:

    70             raise TypeError, "can't pickle %s objects" % base.__name__

    71         state = base(self)

    72     args = (self.__class__, base, state)

    73     try:

    74         getstate = self.__getstate__

    75     except AttributeError:

    76         if getattr(self, "__slots__", None):

    77             raise TypeError("a class that defines __slots__ without "

    78                             "defining __getstate__ cannot be pickled")

    79         try:

    80             dict = self.__dict__

    81         except AttributeError:

    82             dict = None

    83     else:

    84         dict = getstate()

    85     if dict:

    86         return _reconstructor, args, dict

    87     else:

    88         return _reconstructor, args

    89 

    90 # Helper for __reduce_ex__ protocol 2

    91 

    92 def __newobj__(cls, *args):

    93     return cls.__new__(cls, *args)

    94 

    95 def _slotnames(cls):

    96     """Return a list of slot names for a given class.

    97 

    98     This needs to find slots defined by the class and its bases, so we

    99     can't simply return the __slots__ attribute.  We must walk down

   100     the Method Resolution Order and concatenate the __slots__ of each

   101     class found there.  (This assumes classes don't modify their

   102     __slots__ attribute to misrepresent their slots after the class is

   103     defined.)

   104     """

   105 

   106     # Get the value from a cache in the class if possible

   107     names = cls.__dict__.get("__slotnames__")

   108     if names is not None:

   109         return names

   110 

   111     # Not cached -- calculate the value

   112     names = []

   113     if not hasattr(cls, "__slots__"):

   114         # This class has no slots

   115         pass

   116     else:

   117         # Slots found -- gather slot names from all base classes

   118         for c in cls.__mro__:

   119             if "__slots__" in c.__dict__:

   120                 slots = c.__dict__['__slots__']

   121                 # if class has a single slot, it can be given as a string

   122                 if isinstance(slots, basestring):

   123                     slots = (slots,)

   124                 for name in slots:

   125                     # special descriptors

   126                     if name in ("__dict__", "__weakref__"):

   127                         continue

   128                     # mangled names

   129                     elif name.startswith('__') and not name.endswith('__'):

   130                         names.append('_%s%s' % (c.__name__, name))

   131                     else:

   132                         names.append(name)

   133 

   134     # Cache the outcome in the class if at all possible

   135     try:

   136         cls.__slotnames__ = names

   137     except:

   138         pass # But don't die if we can't

   139 

   140     return names

   141 

   142 # A registry of extension codes.  This is an ad-hoc compression

   143 # mechanism.  Whenever a global reference to <module>, <name> is about

   144 # to be pickled, the (<module>, <name>) tuple is looked up here to see

   145 # if it is a registered extension code for it.  Extension codes are

   146 # universal, so that the meaning of a pickle does not depend on

   147 # context.  (There are also some codes reserved for local use that

   148 # don't have this restriction.)  Codes are positive ints; 0 is

   149 # reserved.

   150 

   151 _extension_registry = {}                # key -> code

   152 _inverted_registry = {}                 # code -> key

   153 _extension_cache = {}                   # code -> object

   154 # Don't ever rebind those names:  cPickle grabs a reference to them when

   155 # it's initialized, and won't see a rebinding.

   156 

   157 def add_extension(module, name, code):

   158     """Register an extension code."""

   159     code = int(code)

   160     if not 1 <= code <= 0x7fffffff:

   161         raise ValueError, "code out of range"

   162     key = (module, name)

   163     if (_extension_registry.get(key) == code and

   164         _inverted_registry.get(code) == key):

   165         return # Redundant registrations are benign

   166     if key in _extension_registry:

   167         raise ValueError("key %s is already registered with code %s" %

   168                          (key, _extension_registry[key]))

   169     if code in _inverted_registry:

   170         raise ValueError("code %s is already in use for key %s" %

   171                          (code, _inverted_registry[code]))

   172     _extension_registry[key] = code

   173     _inverted_registry[code] = key

   174 

   175 def remove_extension(module, name, code):

   176     """Unregister an extension code.  For testing only."""

   177     key = (module, name)

   178     if (_extension_registry.get(key) != code or

   179         _inverted_registry.get(code) != key):

   180         raise ValueError("key %s is not registered with code %s" %

   181                          (key, code))

   182     del _extension_registry[key]

   183     del _inverted_registry[code]

   184     if code in _extension_cache:

   185         del _extension_cache[code]

   186 

   187 def clear_extension_cache():

   188     _extension_cache.clear()

   189 

   190 # Standard extension code assignments

   191 

   192 # Reserved ranges

   193 

   194 # First  Last Count  Purpose

   195 #     1   127   127  Reserved for Python standard library

   196 #   128   191    64  Reserved for Zope

   197 #   192   239    48  Reserved for 3rd parties

   198 #   240   255    16  Reserved for private use (will never be assigned)

   199 #   256   Inf   Inf  Reserved for future assignment

   200 

   201 # Extension codes are assigned by the Python Software Foundation.