456 lines
13 KiB
Python
456 lines
13 KiB
Python
# encoding: utf-8
|
|
"""Pickle related utilities. Perhaps this should be called 'can'."""
|
|
|
|
# Copyright (c) IPython Development Team.
|
|
# Distributed under the terms of the Modified BSD License.
|
|
|
|
import warnings
|
|
warnings.warn("ipykernel.pickleutil is deprecated. It has moved to ipyparallel.", DeprecationWarning)
|
|
|
|
import copy
|
|
import sys
|
|
from types import FunctionType
|
|
|
|
try:
|
|
import cPickle as pickle
|
|
except ImportError:
|
|
import pickle
|
|
|
|
from ipython_genutils import py3compat
|
|
from ipython_genutils.importstring import import_item
|
|
from ipython_genutils.py3compat import string_types, iteritems, buffer_to_bytes, buffer_to_bytes_py2
|
|
|
|
# This registers a hook when it's imported
|
|
try:
|
|
# available since ipyparallel 5.1.1
|
|
from ipyparallel.serialize import codeutil
|
|
except ImportError:
|
|
# Deprecated since ipykernel 4.3.1
|
|
from ipykernel import codeutil
|
|
|
|
from traitlets.log import get_logger
|
|
|
|
if py3compat.PY3:
|
|
buffer = memoryview
|
|
class_type = type
|
|
else:
|
|
from types import ClassType
|
|
class_type = (type, ClassType)
|
|
|
|
try:
|
|
PICKLE_PROTOCOL = pickle.DEFAULT_PROTOCOL
|
|
except AttributeError:
|
|
PICKLE_PROTOCOL = pickle.HIGHEST_PROTOCOL
|
|
|
|
def _get_cell_type(a=None):
|
|
"""the type of a closure cell doesn't seem to be importable,
|
|
so just create one
|
|
"""
|
|
def inner():
|
|
return a
|
|
return type(py3compat.get_closure(inner)[0])
|
|
|
|
cell_type = _get_cell_type()
|
|
|
|
#-------------------------------------------------------------------------------
|
|
# Functions
|
|
#-------------------------------------------------------------------------------
|
|
|
|
|
|
def interactive(f):
|
|
"""decorator for making functions appear as interactively defined.
|
|
This results in the function being linked to the user_ns as globals()
|
|
instead of the module globals().
|
|
"""
|
|
|
|
# build new FunctionType, so it can have the right globals
|
|
# interactive functions never have closures, that's kind of the point
|
|
if isinstance(f, FunctionType):
|
|
mainmod = __import__('__main__')
|
|
f = FunctionType(f.__code__, mainmod.__dict__,
|
|
f.__name__, f.__defaults__,
|
|
)
|
|
# associate with __main__ for uncanning
|
|
f.__module__ = '__main__'
|
|
return f
|
|
|
|
|
|
def use_dill():
|
|
"""use dill to expand serialization support
|
|
|
|
adds support for object methods and closures to serialization.
|
|
"""
|
|
# import dill causes most of the magic
|
|
import dill
|
|
|
|
# dill doesn't work with cPickle,
|
|
# tell the two relevant modules to use plain pickle
|
|
|
|
global pickle
|
|
pickle = dill
|
|
|
|
try:
|
|
from ipykernel import serialize
|
|
except ImportError:
|
|
pass
|
|
else:
|
|
serialize.pickle = dill
|
|
|
|
# disable special function handling, let dill take care of it
|
|
can_map.pop(FunctionType, None)
|
|
|
|
def use_cloudpickle():
|
|
"""use cloudpickle to expand serialization support
|
|
|
|
adds support for object methods and closures to serialization.
|
|
"""
|
|
import cloudpickle
|
|
|
|
global pickle
|
|
pickle = cloudpickle
|
|
|
|
try:
|
|
from ipykernel import serialize
|
|
except ImportError:
|
|
pass
|
|
else:
|
|
serialize.pickle = cloudpickle
|
|
|
|
# disable special function handling, let cloudpickle take care of it
|
|
can_map.pop(FunctionType, None)
|
|
|
|
|
|
#-------------------------------------------------------------------------------
|
|
# Classes
|
|
#-------------------------------------------------------------------------------
|
|
|
|
|
|
class CannedObject(object):
|
|
def __init__(self, obj, keys=[], hook=None):
|
|
"""can an object for safe pickling
|
|
|
|
Parameters
|
|
==========
|
|
|
|
obj:
|
|
The object to be canned
|
|
keys: list (optional)
|
|
list of attribute names that will be explicitly canned / uncanned
|
|
hook: callable (optional)
|
|
An optional extra callable,
|
|
which can do additional processing of the uncanned object.
|
|
|
|
large data may be offloaded into the buffers list,
|
|
used for zero-copy transfers.
|
|
"""
|
|
self.keys = keys
|
|
self.obj = copy.copy(obj)
|
|
self.hook = can(hook)
|
|
for key in keys:
|
|
setattr(self.obj, key, can(getattr(obj, key)))
|
|
|
|
self.buffers = []
|
|
|
|
def get_object(self, g=None):
|
|
if g is None:
|
|
g = {}
|
|
obj = self.obj
|
|
for key in self.keys:
|
|
setattr(obj, key, uncan(getattr(obj, key), g))
|
|
|
|
if self.hook:
|
|
self.hook = uncan(self.hook, g)
|
|
self.hook(obj, g)
|
|
return self.obj
|
|
|
|
|
|
class Reference(CannedObject):
|
|
"""object for wrapping a remote reference by name."""
|
|
def __init__(self, name):
|
|
if not isinstance(name, string_types):
|
|
raise TypeError("illegal name: %r"%name)
|
|
self.name = name
|
|
self.buffers = []
|
|
|
|
def __repr__(self):
|
|
return "<Reference: %r>"%self.name
|
|
|
|
def get_object(self, g=None):
|
|
if g is None:
|
|
g = {}
|
|
|
|
return eval(self.name, g)
|
|
|
|
|
|
class CannedCell(CannedObject):
|
|
"""Can a closure cell"""
|
|
def __init__(self, cell):
|
|
self.cell_contents = can(cell.cell_contents)
|
|
|
|
def get_object(self, g=None):
|
|
cell_contents = uncan(self.cell_contents, g)
|
|
def inner():
|
|
return cell_contents
|
|
return py3compat.get_closure(inner)[0]
|
|
|
|
|
|
class CannedFunction(CannedObject):
|
|
|
|
def __init__(self, f):
|
|
self._check_type(f)
|
|
self.code = f.__code__
|
|
if f.__defaults__:
|
|
self.defaults = [ can(fd) for fd in f.__defaults__ ]
|
|
else:
|
|
self.defaults = None
|
|
|
|
closure = py3compat.get_closure(f)
|
|
if closure:
|
|
self.closure = tuple( can(cell) for cell in closure )
|
|
else:
|
|
self.closure = None
|
|
|
|
self.module = f.__module__ or '__main__'
|
|
self.__name__ = f.__name__
|
|
self.buffers = []
|
|
|
|
def _check_type(self, obj):
|
|
assert isinstance(obj, FunctionType), "Not a function type"
|
|
|
|
def get_object(self, g=None):
|
|
# try to load function back into its module:
|
|
if not self.module.startswith('__'):
|
|
__import__(self.module)
|
|
g = sys.modules[self.module].__dict__
|
|
|
|
if g is None:
|
|
g = {}
|
|
if self.defaults:
|
|
defaults = tuple(uncan(cfd, g) for cfd in self.defaults)
|
|
else:
|
|
defaults = None
|
|
if self.closure:
|
|
closure = tuple(uncan(cell, g) for cell in self.closure)
|
|
else:
|
|
closure = None
|
|
newFunc = FunctionType(self.code, g, self.__name__, defaults, closure)
|
|
return newFunc
|
|
|
|
class CannedClass(CannedObject):
|
|
|
|
def __init__(self, cls):
|
|
self._check_type(cls)
|
|
self.name = cls.__name__
|
|
self.old_style = not isinstance(cls, type)
|
|
self._canned_dict = {}
|
|
for k,v in cls.__dict__.items():
|
|
if k not in ('__weakref__', '__dict__'):
|
|
self._canned_dict[k] = can(v)
|
|
if self.old_style:
|
|
mro = []
|
|
else:
|
|
mro = cls.mro()
|
|
|
|
self.parents = [ can(c) for c in mro[1:] ]
|
|
self.buffers = []
|
|
|
|
def _check_type(self, obj):
|
|
assert isinstance(obj, class_type), "Not a class type"
|
|
|
|
def get_object(self, g=None):
|
|
parents = tuple(uncan(p, g) for p in self.parents)
|
|
return type(self.name, parents, uncan_dict(self._canned_dict, g=g))
|
|
|
|
class CannedArray(CannedObject):
|
|
def __init__(self, obj):
|
|
from numpy import ascontiguousarray
|
|
self.shape = obj.shape
|
|
self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
|
|
self.pickled = False
|
|
if sum(obj.shape) == 0:
|
|
self.pickled = True
|
|
elif obj.dtype == 'O':
|
|
# can't handle object dtype with buffer approach
|
|
self.pickled = True
|
|
elif obj.dtype.fields and any(dt == 'O' for dt,sz in obj.dtype.fields.values()):
|
|
self.pickled = True
|
|
if self.pickled:
|
|
# just pickle it
|
|
self.buffers = [pickle.dumps(obj, PICKLE_PROTOCOL)]
|
|
else:
|
|
# ensure contiguous
|
|
obj = ascontiguousarray(obj, dtype=None)
|
|
self.buffers = [buffer(obj)]
|
|
|
|
def get_object(self, g=None):
|
|
from numpy import frombuffer
|
|
data = self.buffers[0]
|
|
if self.pickled:
|
|
# we just pickled it
|
|
return pickle.loads(buffer_to_bytes_py2(data))
|
|
else:
|
|
if not py3compat.PY3 and isinstance(data, memoryview):
|
|
# frombuffer doesn't accept memoryviews on Python 2,
|
|
# so cast to old-style buffer
|
|
data = buffer(data.tobytes())
|
|
return frombuffer(data, dtype=self.dtype).reshape(self.shape)
|
|
|
|
|
|
class CannedBytes(CannedObject):
|
|
wrap = staticmethod(buffer_to_bytes)
|
|
|
|
def __init__(self, obj):
|
|
self.buffers = [obj]
|
|
|
|
def get_object(self, g=None):
|
|
data = self.buffers[0]
|
|
return self.wrap(data)
|
|
|
|
class CannedBuffer(CannedBytes):
|
|
wrap = buffer
|
|
|
|
class CannedMemoryView(CannedBytes):
|
|
wrap = memoryview
|
|
|
|
#-------------------------------------------------------------------------------
|
|
# Functions
|
|
#-------------------------------------------------------------------------------
|
|
|
|
def _import_mapping(mapping, original=None):
|
|
"""import any string-keys in a type mapping
|
|
|
|
"""
|
|
log = get_logger()
|
|
log.debug("Importing canning map")
|
|
for key,value in list(mapping.items()):
|
|
if isinstance(key, string_types):
|
|
try:
|
|
cls = import_item(key)
|
|
except Exception:
|
|
if original and key not in original:
|
|
# only message on user-added classes
|
|
log.error("canning class not importable: %r", key, exc_info=True)
|
|
mapping.pop(key)
|
|
else:
|
|
mapping[cls] = mapping.pop(key)
|
|
|
|
def istype(obj, check):
|
|
"""like isinstance(obj, check), but strict
|
|
|
|
This won't catch subclasses.
|
|
"""
|
|
if isinstance(check, tuple):
|
|
for cls in check:
|
|
if type(obj) is cls:
|
|
return True
|
|
return False
|
|
else:
|
|
return type(obj) is check
|
|
|
|
def can(obj):
|
|
"""prepare an object for pickling"""
|
|
|
|
import_needed = False
|
|
|
|
for cls,canner in iteritems(can_map):
|
|
if isinstance(cls, string_types):
|
|
import_needed = True
|
|
break
|
|
elif istype(obj, cls):
|
|
return canner(obj)
|
|
|
|
if import_needed:
|
|
# perform can_map imports, then try again
|
|
# this will usually only happen once
|
|
_import_mapping(can_map, _original_can_map)
|
|
return can(obj)
|
|
|
|
return obj
|
|
|
|
def can_class(obj):
|
|
if isinstance(obj, class_type) and obj.__module__ == '__main__':
|
|
return CannedClass(obj)
|
|
else:
|
|
return obj
|
|
|
|
def can_dict(obj):
|
|
"""can the *values* of a dict"""
|
|
if istype(obj, dict):
|
|
newobj = {}
|
|
for k, v in iteritems(obj):
|
|
newobj[k] = can(v)
|
|
return newobj
|
|
else:
|
|
return obj
|
|
|
|
sequence_types = (list, tuple, set)
|
|
|
|
def can_sequence(obj):
|
|
"""can the elements of a sequence"""
|
|
if istype(obj, sequence_types):
|
|
t = type(obj)
|
|
return t([can(i) for i in obj])
|
|
else:
|
|
return obj
|
|
|
|
def uncan(obj, g=None):
|
|
"""invert canning"""
|
|
|
|
import_needed = False
|
|
for cls,uncanner in iteritems(uncan_map):
|
|
if isinstance(cls, string_types):
|
|
import_needed = True
|
|
break
|
|
elif isinstance(obj, cls):
|
|
return uncanner(obj, g)
|
|
|
|
if import_needed:
|
|
# perform uncan_map imports, then try again
|
|
# this will usually only happen once
|
|
_import_mapping(uncan_map, _original_uncan_map)
|
|
return uncan(obj, g)
|
|
|
|
return obj
|
|
|
|
def uncan_dict(obj, g=None):
|
|
if istype(obj, dict):
|
|
newobj = {}
|
|
for k, v in iteritems(obj):
|
|
newobj[k] = uncan(v,g)
|
|
return newobj
|
|
else:
|
|
return obj
|
|
|
|
def uncan_sequence(obj, g=None):
|
|
if istype(obj, sequence_types):
|
|
t = type(obj)
|
|
return t([uncan(i,g) for i in obj])
|
|
else:
|
|
return obj
|
|
|
|
#-------------------------------------------------------------------------------
|
|
# API dictionaries
|
|
#-------------------------------------------------------------------------------
|
|
|
|
# These dicts can be extended for custom serialization of new objects
|
|
|
|
can_map = {
|
|
'numpy.ndarray' : CannedArray,
|
|
FunctionType : CannedFunction,
|
|
bytes : CannedBytes,
|
|
memoryview : CannedMemoryView,
|
|
cell_type : CannedCell,
|
|
class_type : can_class,
|
|
}
|
|
if buffer is not memoryview:
|
|
can_map[buffer] = CannedBuffer
|
|
|
|
uncan_map = {
|
|
CannedObject : lambda obj, g: obj.get_object(g),
|
|
dict : uncan_dict,
|
|
}
|
|
|
|
# for use in _import_mapping:
|
|
_original_can_map = can_map.copy()
|
|
_original_uncan_map = uncan_map.copy()
|