update pickleutil imports

ipython_kernel/pickleutil.py

@@ -1,425 +1,426 b''
 # 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 copy
 import logging
 import sys
 from types import FunctionType
 
 try:
     import cPickle as pickle
 except ImportError:
     import pickle
 
+from IPython.utils import py3compat
+from IPython.utils.importstring import import_item
+from IPython.utils.py3compat import string_types, iteritems, buffer_to_bytes_py2
+
 from . import codeutil  # This registers a hook when it's imported
-from . import py3compat
-from .importstring import import_item
-from .py3compat import string_types, iteritems, buffer_to_bytes_py2
 
 from IPython.config import Application
 from IPython.utils.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 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 IPython.kernel.zmq 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.
     """
     from cloud.serialization import cloudpickle
     
     global pickle
     pickle = cloudpickle
 
     try:
         from IPython.kernel.zmq 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:
             return frombuffer(data, dtype=self.dtype).reshape(self.shape)
 
 
 class CannedBytes(CannedObject):
     wrap = bytes
     def __init__(self, obj):
         self.buffers = [obj]
     
     def get_object(self, g=None):
         data = self.buffers[0]
         return self.wrap(data)
 
 def CannedBuffer(CannedBytes):
     wrap = buffer
 
 #-------------------------------------------------------------------------------
 # 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
 
 def _uncan_dependent_hook(dep, g=None):
     dep.check_dependency()
     
 def can_dependent(obj):
     return CannedObject(obj, keys=('f', 'df'), hook=_uncan_dependent_hook)
 
 #-------------------------------------------------------------------------------
 # API dictionaries
 #-------------------------------------------------------------------------------
 
 # These dicts can be extended for custom serialization of new objects
 
 can_map = {
     'IPython.parallel.dependent' : can_dependent,
     'numpy.ndarray' : CannedArray,
     FunctionType : CannedFunction,
     bytes : CannedBytes,
     buffer : CannedBuffer,
     cell_type : CannedCell,
     class_type : can_class,
 }
 
 uncan_map = {
     CannedObject : lambda obj, g: obj.get_object(g),
 }
 
 # for use in _import_mapping:
 _original_can_map = can_map.copy()
 _original_uncan_map = uncan_map.copy()

ipython_kernel/serialize.py

@@ -1,179 +1,179 b''
 """serialization utilities for apply messages"""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 try:
     import cPickle
     pickle = cPickle
 except:
     cPickle = None
     import pickle
 
 # IPython imports
 from IPython.utils.py3compat import PY3, buffer_to_bytes_py2
 from IPython.utils.data import flatten
-from IPython.utils.pickleutil import (
+from ipython_kernel.pickleutil import (
     can, uncan, can_sequence, uncan_sequence, CannedObject,
     istype, sequence_types, PICKLE_PROTOCOL,
 )
 from jupyter_client.session import MAX_ITEMS, MAX_BYTES
 
 
 if PY3:
     buffer = memoryview
 
 #-----------------------------------------------------------------------------
 # Serialization Functions
 #-----------------------------------------------------------------------------
 
 
 def _extract_buffers(obj, threshold=MAX_BYTES):
     """extract buffers larger than a certain threshold"""
     buffers = []
     if isinstance(obj, CannedObject) and obj.buffers:
         for i,buf in enumerate(obj.buffers):
             if len(buf) > threshold:
                 # buffer larger than threshold, prevent pickling
                 obj.buffers[i] = None
                 buffers.append(buf)
             elif isinstance(buf, buffer):
                 # buffer too small for separate send, coerce to bytes
                 # because pickling buffer objects just results in broken pointers
                 obj.buffers[i] = bytes(buf)
     return buffers
 
 def _restore_buffers(obj, buffers):
     """restore buffers extracted by """
     if isinstance(obj, CannedObject) and obj.buffers:
         for i,buf in enumerate(obj.buffers):
             if buf is None:
                 obj.buffers[i] = buffers.pop(0)
 
 def serialize_object(obj, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS):
     """Serialize an object into a list of sendable buffers.
 
     Parameters
     ----------
 
     obj : object
         The object to be serialized
     buffer_threshold : int
         The threshold (in bytes) for pulling out data buffers
         to avoid pickling them.
     item_threshold : int
         The maximum number of items over which canning will iterate.
         Containers (lists, dicts) larger than this will be pickled without
         introspection.
 
     Returns
     -------
     [bufs] : list of buffers representing the serialized object.
     """
     buffers = []
     if istype(obj, sequence_types) and len(obj) < item_threshold:
         cobj = can_sequence(obj)
         for c in cobj:
             buffers.extend(_extract_buffers(c, buffer_threshold))
     elif istype(obj, dict) and len(obj) < item_threshold:
         cobj = {}
         for k in sorted(obj):
             c = can(obj[k])
             buffers.extend(_extract_buffers(c, buffer_threshold))
             cobj[k] = c
     else:
         cobj = can(obj)
         buffers.extend(_extract_buffers(cobj, buffer_threshold))
 
     buffers.insert(0, pickle.dumps(cobj, PICKLE_PROTOCOL))
     return buffers
 
 def deserialize_object(buffers, g=None):
     """reconstruct an object serialized by serialize_object from data buffers.
 
     Parameters
     ----------
 
     bufs : list of buffers/bytes
 
     g : globals to be used when uncanning
 
     Returns
     -------
 
     (newobj, bufs) : unpacked object, and the list of remaining unused buffers.
     """
     bufs = list(buffers)
     pobj = buffer_to_bytes_py2(bufs.pop(0))
     canned = pickle.loads(pobj)
     if istype(canned, sequence_types) and len(canned) < MAX_ITEMS:
         for c in canned:
             _restore_buffers(c, bufs)
         newobj = uncan_sequence(canned, g)
     elif istype(canned, dict) and len(canned) < MAX_ITEMS:
         newobj = {}
         for k in sorted(canned):
             c = canned[k]
             _restore_buffers(c, bufs)
             newobj[k] = uncan(c, g)
     else:
         _restore_buffers(canned, bufs)
         newobj = uncan(canned, g)
 
     return newobj, bufs
 
 def pack_apply_message(f, args, kwargs, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS):
     """pack up a function, args, and kwargs to be sent over the wire
 
     Each element of args/kwargs will be canned for special treatment,
     but inspection will not go any deeper than that.
 
     Any object whose data is larger than `threshold`  will not have their data copied
     (only numpy arrays and bytes/buffers support zero-copy)
 
     Message will be a list of bytes/buffers of the format:
 
     [ cf, pinfo, <arg_bufs>, <kwarg_bufs> ]
 
     With length at least two + len(args) + len(kwargs)
     """
 
     arg_bufs = flatten(serialize_object(arg, buffer_threshold, item_threshold) for arg in args)
 
     kw_keys = sorted(kwargs.keys())
     kwarg_bufs = flatten(serialize_object(kwargs[key], buffer_threshold, item_threshold) for key in kw_keys)
 
     info = dict(nargs=len(args), narg_bufs=len(arg_bufs), kw_keys=kw_keys)
 
     msg = [pickle.dumps(can(f), PICKLE_PROTOCOL)]
     msg.append(pickle.dumps(info, PICKLE_PROTOCOL))
     msg.extend(arg_bufs)
     msg.extend(kwarg_bufs)
 
     return msg
 
 def unpack_apply_message(bufs, g=None, copy=True):
     """unpack f,args,kwargs from buffers packed by pack_apply_message()
     Returns: original f,args,kwargs"""
     bufs = list(bufs) # allow us to pop
     assert len(bufs) >= 2, "not enough buffers!"
     pf = buffer_to_bytes_py2(bufs.pop(0))
     f = uncan(pickle.loads(pf), g)
     pinfo = buffer_to_bytes_py2(bufs.pop(0))
     info = pickle.loads(pinfo)
     arg_bufs, kwarg_bufs = bufs[:info['narg_bufs']], bufs[info['narg_bufs']:]
 
     args = []
     for i in range(info['nargs']):
         arg, arg_bufs = deserialize_object(arg_bufs, g)
         args.append(arg)
     args = tuple(args)
     assert not arg_bufs, "Shouldn't be any arg bufs left over"
 
     kwargs = {}
     for key in info['kw_keys']:
         kwarg, kwarg_bufs = deserialize_object(kwarg_bufs, g)
         kwargs[key] = kwarg
     assert not kwarg_bufs, "Shouldn't be any kwarg bufs left over"
 
     return f,args,kwargs

ipython_kernel/tests/test_pickleutil.py

@@ -1,62 +1,62 b''
 
 import pickle
 
 import nose.tools as nt
-from IPython.utils import codeutil
-from IPython.utils.pickleutil import can, uncan
+from ipython_kernel import codeutil
+from ipython_kernel.pickleutil import can, uncan
 
 def interactive(f):
     f.__module__ = '__main__'
     return f
 
 def dumps(obj):
     return pickle.dumps(can(obj))
 
 def loads(obj):
     return uncan(pickle.loads(obj))
 
 def test_no_closure():
     @interactive
     def foo():
         a = 5
         return a
     
     pfoo = dumps(foo)
     bar = loads(pfoo)
     nt.assert_equal(foo(), bar())
 
 def test_generator_closure():
     # this only creates a closure on Python 3
     @interactive
     def foo():
         i = 'i'
         r = [ i for j in (1,2) ]
         return r
     
     pfoo = dumps(foo)
     bar = loads(pfoo)
     nt.assert_equal(foo(), bar())
 
 def test_nested_closure():
     @interactive
     def foo():
         i = 'i'
         def g():
             return i
         return g()
     
     pfoo = dumps(foo)
     bar = loads(pfoo)
     nt.assert_equal(foo(), bar())
 
 def test_closure():
     i = 'i'
     @interactive
     def foo():
         return i
     
     pfoo = dumps(foo)
     bar = loads(pfoo)
     nt.assert_equal(foo(), bar())
 
     
  No newline at end of file

ipython_kernel/tests/test_serialize.py

@@ -1,208 +1,208 b''
 """test serialization tools"""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import pickle
 from collections import namedtuple
 
 import nose.tools as nt
 
 # from unittest import TestCaes
 from ipython_kernel.serialize import serialize_object, deserialize_object
 from IPython.testing import decorators as dec
-from IPython.utils.pickleutil import CannedArray, CannedClass
+from ipython_kernel.pickleutil import CannedArray, CannedClass
 from IPython.utils.py3compat import iteritems
 from IPython.parallel import interactive
 
 #-------------------------------------------------------------------------------
 # Globals and Utilities
 #-------------------------------------------------------------------------------
 
 def roundtrip(obj):
     """roundtrip an object through serialization"""
     bufs = serialize_object(obj)
     obj2, remainder = deserialize_object(bufs)
     nt.assert_equals(remainder, [])
     return obj2
 
 class C(object):
     """dummy class for """
 
     def __init__(self, **kwargs):
         for key,value in iteritems(kwargs):
             setattr(self, key, value)
 
 SHAPES = ((100,), (1024,10), (10,8,6,5), (), (0,))
 DTYPES = ('uint8', 'float64', 'int32', [('g', 'float32')], '|S10')
 
 #-------------------------------------------------------------------------------
 # Tests
 #-------------------------------------------------------------------------------
 
 def new_array(shape, dtype):
     import numpy
     return numpy.random.random(shape).astype(dtype)
 
 def test_roundtrip_simple():
     for obj in [
         'hello',
         dict(a='b', b=10),
         [1,2,'hi'],
         (b'123', 'hello'),
     ]:
         obj2 = roundtrip(obj)
         nt.assert_equal(obj, obj2)
 
 def test_roundtrip_nested():
     for obj in [
         dict(a=range(5), b={1:b'hello'}),
         [range(5),[range(3),(1,[b'whoda'])]],
     ]:
         obj2 = roundtrip(obj)
         nt.assert_equal(obj, obj2)
 
 def test_roundtrip_buffered():
     for obj in [
         dict(a=b"x"*1025),
         b"hello"*500,
         [b"hello"*501, 1,2,3]
     ]:
         bufs = serialize_object(obj)
         nt.assert_equal(len(bufs), 2)
         obj2, remainder = deserialize_object(bufs)
         nt.assert_equal(remainder, [])
         nt.assert_equal(obj, obj2)
 
 @dec.skip_without('numpy')
 def test_numpy():
     import numpy
     from numpy.testing.utils import assert_array_equal
     for shape in SHAPES:
         for dtype in DTYPES:
             A = new_array(shape, dtype=dtype)
             bufs = serialize_object(A)
             B, r = deserialize_object(bufs)
             nt.assert_equal(r, [])
             nt.assert_equal(A.shape, B.shape)
             nt.assert_equal(A.dtype, B.dtype)
             assert_array_equal(A,B)
 
 @dec.skip_without('numpy')
 def test_recarray():
     import numpy
     from numpy.testing.utils import assert_array_equal
     for shape in SHAPES:
         for dtype in [
             [('f', float), ('s', '|S10')],
             [('n', int), ('s', '|S1'), ('u', 'uint32')],
         ]:
             A = new_array(shape, dtype=dtype)
 
             bufs = serialize_object(A)
             B, r = deserialize_object(bufs)
             nt.assert_equal(r, [])
             nt.assert_equal(A.shape, B.shape)
             nt.assert_equal(A.dtype, B.dtype)
             assert_array_equal(A,B)
 
 @dec.skip_without('numpy')
 def test_numpy_in_seq():
     import numpy
     from numpy.testing.utils import assert_array_equal
     for shape in SHAPES:
         for dtype in DTYPES:
             A = new_array(shape, dtype=dtype)
             bufs = serialize_object((A,1,2,b'hello'))
             canned = pickle.loads(bufs[0])
             nt.assert_is_instance(canned[0], CannedArray)
             tup, r = deserialize_object(bufs)
             B = tup[0]
             nt.assert_equal(r, [])
             nt.assert_equal(A.shape, B.shape)
             nt.assert_equal(A.dtype, B.dtype)
             assert_array_equal(A,B)
 
 @dec.skip_without('numpy')
 def test_numpy_in_dict():
     import numpy
     from numpy.testing.utils import assert_array_equal
     for shape in SHAPES:
         for dtype in DTYPES:
             A = new_array(shape, dtype=dtype)
             bufs = serialize_object(dict(a=A,b=1,c=range(20)))
             canned = pickle.loads(bufs[0])
             nt.assert_is_instance(canned['a'], CannedArray)
             d, r = deserialize_object(bufs)
             B = d['a']
             nt.assert_equal(r, [])
             nt.assert_equal(A.shape, B.shape)
             nt.assert_equal(A.dtype, B.dtype)
             assert_array_equal(A,B)
 
 def test_class():
     @interactive
     class C(object):
         a=5
     bufs = serialize_object(dict(C=C))
     canned = pickle.loads(bufs[0])
     nt.assert_is_instance(canned['C'], CannedClass)
     d, r = deserialize_object(bufs)
     C2 = d['C']
     nt.assert_equal(C2.a, C.a)
 
 def test_class_oldstyle():
     @interactive
     class C:
         a=5
 
     bufs = serialize_object(dict(C=C))
     canned = pickle.loads(bufs[0])
     nt.assert_is_instance(canned['C'], CannedClass)
     d, r = deserialize_object(bufs)
     C2 = d['C']
     nt.assert_equal(C2.a, C.a)
 
 def test_tuple():
     tup = (lambda x:x, 1)
     bufs = serialize_object(tup)
     canned = pickle.loads(bufs[0])
     nt.assert_is_instance(canned, tuple)
     t2, r = deserialize_object(bufs)
     nt.assert_equal(t2[0](t2[1]), tup[0](tup[1]))
 
 point = namedtuple('point', 'x y')
 
 def test_namedtuple():
     p = point(1,2)
     bufs = serialize_object(p)
     canned = pickle.loads(bufs[0])
     nt.assert_is_instance(canned, point)
     p2, r = deserialize_object(bufs, globals())
     nt.assert_equal(p2.x, p.x)
     nt.assert_equal(p2.y, p.y)
 
 def test_list():
     lis = [lambda x:x, 1]
     bufs = serialize_object(lis)
     canned = pickle.loads(bufs[0])
     nt.assert_is_instance(canned, list)
     l2, r = deserialize_object(bufs)
     nt.assert_equal(l2[0](l2[1]), lis[0](lis[1]))
 
 def test_class_inheritance():
     @interactive
     class C(object):
         a=5
 
     @interactive
     class D(C):
         b=10
 
     bufs = serialize_object(dict(D=D))
     canned = pickle.loads(bufs[0])
     nt.assert_is_instance(canned['D'], CannedClass)
     d, r = deserialize_object(bufs)
     D2 = d['D']
     nt.assert_equal(D2.a, D.a)
     nt.assert_equal(D2.b, D.b)

ipython_parallel/__init__.py

@@ -1,72 +1,72 b''
 """The IPython ZMQ-based parallel computing interface.
 
 Authors:
 
 * MinRK
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2011 The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import os
 import warnings
 
 import zmq
 
 from IPython.config.configurable import MultipleInstanceError
 from IPython.utils.zmqrelated import check_for_zmq
 
 min_pyzmq = '2.1.11'
 
 check_for_zmq(min_pyzmq, 'ipython_parallel')
 
-from IPython.utils.pickleutil import Reference
+from ipython_kernel.pickleutil import Reference
 
 from .client.asyncresult import *
 from .client.client import Client
 from .client.remotefunction import *
 from .client.view import *
 from .controller.dependency import *
 from .error import *
 from .util import interactive
 
 #-----------------------------------------------------------------------------
 # Functions
 #-----------------------------------------------------------------------------
 
 
 def bind_kernel(**kwargs):
     """Bind an Engine's Kernel to be used as a full IPython kernel.
     
     This allows a running Engine to be used simultaneously as a full IPython kernel
     with the QtConsole or other frontends.
     
     This function returns immediately.
     """
     from IPython.kernel.zmq.kernelapp import IPKernelApp
     from ipython_parallel.apps.ipengineapp import IPEngineApp
     
     # first check for IPKernelApp, in which case this should be a no-op
     # because there is already a bound kernel
     if IPKernelApp.initialized() and isinstance(IPKernelApp._instance, IPKernelApp):
         return
     
     if IPEngineApp.initialized():
         try:
             app = IPEngineApp.instance()
         except MultipleInstanceError:
             pass
         else:
             return app.bind_kernel(**kwargs)
     
     raise RuntimeError("bind_kernel be called from an IPEngineApp instance")
     
 
 

ipython_parallel/client/view.py

@@ -1,1125 +1,1125 b''
 """Views of remote engines."""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 from __future__ import print_function
 
 import imp
 import sys
 import warnings
 from contextlib import contextmanager
 from types import ModuleType
 
 import zmq
 
 from IPython.testing.skipdoctest import skip_doctest
 from IPython.utils import pickleutil
 from IPython.utils.traitlets import (
     HasTraits, Any, Bool, List, Dict, Set, Instance, CFloat, Integer
 )
 from decorator import decorator
 
 from ipython_parallel import util
 from ipython_parallel.controller.dependency import Dependency, dependent
 from IPython.utils.py3compat import string_types, iteritems, PY3
 
 from . import map as Map
 from .asyncresult import AsyncResult, AsyncMapResult
 from .remotefunction import ParallelFunction, parallel, remote, getname
 
 #-----------------------------------------------------------------------------
 # Decorators
 #-----------------------------------------------------------------------------
 
 @decorator
 def save_ids(f, self, *args, **kwargs):
     """Keep our history and outstanding attributes up to date after a method call."""
     n_previous = len(self.client.history)
     try:
         ret = f(self, *args, **kwargs)
     finally:
         nmsgs = len(self.client.history) - n_previous
         msg_ids = self.client.history[-nmsgs:]
         self.history.extend(msg_ids)
         self.outstanding.update(msg_ids)
     return ret
 
 @decorator
 def sync_results(f, self, *args, **kwargs):
     """sync relevant results from self.client to our results attribute."""
     if self._in_sync_results:
         return f(self, *args, **kwargs)
     self._in_sync_results = True
     try:
         ret = f(self, *args, **kwargs)
     finally:
         self._in_sync_results = False
         self._sync_results()
     return ret
 
 @decorator
 def spin_after(f, self, *args, **kwargs):
     """call spin after the method."""
     ret = f(self, *args, **kwargs)
     self.spin()
     return ret
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
 @skip_doctest
 class View(HasTraits):
     """Base View class for more convenint apply(f,*args,**kwargs) syntax via attributes.
 
     Don't use this class, use subclasses.
 
     Methods
     -------
 
     spin
         flushes incoming results and registration state changes
         control methods spin, and requesting `ids` also ensures up to date
 
     wait
         wait on one or more msg_ids
 
     execution methods
         apply
         legacy: execute, run
 
     data movement
         push, pull, scatter, gather
 
     query methods
         get_result, queue_status, purge_results, result_status
 
     control methods
         abort, shutdown
 
     """
     # flags
     block=Bool(False)
     track=Bool(True)
     targets = Any()
 
     history=List()
     outstanding = Set()
     results = Dict()
     client = Instance('ipython_parallel.Client', allow_none=True)
 
     _socket = Instance('zmq.Socket', allow_none=True)
     _flag_names = List(['targets', 'block', 'track'])
     _in_sync_results = Bool(False)
     _targets = Any()
     _idents = Any()
 
     def __init__(self, client=None, socket=None, **flags):
         super(View, self).__init__(client=client, _socket=socket)
         self.results = client.results
         self.block = client.block
 
         self.set_flags(**flags)
 
         assert not self.__class__ is View, "Don't use base View objects, use subclasses"
 
     def __repr__(self):
         strtargets = str(self.targets)
         if len(strtargets) > 16:
             strtargets = strtargets[:12]+'...]'
         return "<%s %s>"%(self.__class__.__name__, strtargets)
     
     def __len__(self):
         if isinstance(self.targets, list):
             return len(self.targets)
         elif isinstance(self.targets, int):
             return 1
         else:
             return len(self.client)
     
     def set_flags(self, **kwargs):
         """set my attribute flags by keyword.
 
         Views determine behavior with a few attributes (`block`, `track`, etc.).
         These attributes can be set all at once by name with this method.
 
         Parameters
         ----------
 
         block : bool
             whether to wait for results
         track : bool
             whether to create a MessageTracker to allow the user to
             safely edit after arrays and buffers during non-copying
             sends.
         """
         for name, value in iteritems(kwargs):
             if name not in self._flag_names:
                 raise KeyError("Invalid name: %r"%name)
             else:
                 setattr(self, name, value)
 
     @contextmanager
     def temp_flags(self, **kwargs):
         """temporarily set flags, for use in `with` statements.
 
         See set_flags for permanent setting of flags
 
         Examples
         --------
 
         >>> view.track=False
         ...
         >>> with view.temp_flags(track=True):
         ...    ar = view.apply(dostuff, my_big_array)
         ...    ar.tracker.wait() # wait for send to finish
         >>> view.track
         False
 
         """
         # preflight: save flags, and set temporaries
         saved_flags = {}
         for f in self._flag_names:
             saved_flags[f] = getattr(self, f)
         self.set_flags(**kwargs)
         # yield to the with-statement block
         try:
             yield
         finally:
             # postflight: restore saved flags
             self.set_flags(**saved_flags)
 
 
     #----------------------------------------------------------------
     # apply
     #----------------------------------------------------------------
 
     def _sync_results(self):
         """to be called by @sync_results decorator
         
         after submitting any tasks.
         """
         delta = self.outstanding.difference(self.client.outstanding)
         completed = self.outstanding.intersection(delta)
         self.outstanding = self.outstanding.difference(completed)
         
     @sync_results
     @save_ids
     def _really_apply(self, f, args, kwargs, block=None, **options):
         """wrapper for client.send_apply_request"""
         raise NotImplementedError("Implement in subclasses")
 
     def apply(self, f, *args, **kwargs):
         """calls ``f(*args, **kwargs)`` on remote engines, returning the result.
 
         This method sets all apply flags via this View's attributes.
 
         Returns :class:`~IPython.parallel.client.asyncresult.AsyncResult`
         instance if ``self.block`` is False, otherwise the return value of
         ``f(*args, **kwargs)``.
         """
         return self._really_apply(f, args, kwargs)
 
     def apply_async(self, f, *args, **kwargs):
         """calls ``f(*args, **kwargs)`` on remote engines in a nonblocking manner.
 
         Returns :class:`~IPython.parallel.client.asyncresult.AsyncResult` instance.
         """
         return self._really_apply(f, args, kwargs, block=False)
 
     @spin_after
     def apply_sync(self, f, *args, **kwargs):
         """calls ``f(*args, **kwargs)`` on remote engines in a blocking manner,
          returning the result.
         """
         return self._really_apply(f, args, kwargs, block=True)
 
     #----------------------------------------------------------------
     # wrappers for client and control methods
     #----------------------------------------------------------------
     @sync_results
     def spin(self):
         """spin the client, and sync"""
         self.client.spin()
 
     @sync_results
     def wait(self, jobs=None, timeout=-1):
         """waits on one or more `jobs`, for up to `timeout` seconds.
 
         Parameters
         ----------
 
         jobs : int, str, or list of ints and/or strs, or one or more AsyncResult objects
                 ints are indices to self.history
                 strs are msg_ids
                 default: wait on all outstanding messages
         timeout : float
                 a time in seconds, after which to give up.
                 default is -1, which means no timeout
 
         Returns
         -------
 
         True : when all msg_ids are done
         False : timeout reached, some msg_ids still outstanding
         """
         if jobs is None:
             jobs = self.history
         return self.client.wait(jobs, timeout)
 
     def abort(self, jobs=None, targets=None, block=None):
         """Abort jobs on my engines.
 
         Parameters
         ----------
 
         jobs : None, str, list of strs, optional
             if None: abort all jobs.
             else: abort specific msg_id(s).
         """
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         jobs = jobs if jobs is not None else list(self.outstanding)
         
         return self.client.abort(jobs=jobs, targets=targets, block=block)
 
     def queue_status(self, targets=None, verbose=False):
         """Fetch the Queue status of my engines"""
         targets = targets if targets is not None else self.targets
         return self.client.queue_status(targets=targets, verbose=verbose)
 
     def purge_results(self, jobs=[], targets=[]):
         """Instruct the controller to forget specific results."""
         if targets is None or targets == 'all':
             targets = self.targets
         return self.client.purge_results(jobs=jobs, targets=targets)
 
     def shutdown(self, targets=None, restart=False, hub=False, block=None):
         """Terminates one or more engine processes, optionally including the hub.
         """
         block = self.block if block is None else block
         if targets is None or targets == 'all':
             targets = self.targets
         return self.client.shutdown(targets=targets, restart=restart, hub=hub, block=block)
 
     @spin_after
     def get_result(self, indices_or_msg_ids=None, block=None, owner=True):
         """return one or more results, specified by history index or msg_id.
 
         See :meth:`IPython.parallel.client.client.Client.get_result` for details.
         """
 
         if indices_or_msg_ids is None:
             indices_or_msg_ids = -1
         if isinstance(indices_or_msg_ids, int):
             indices_or_msg_ids = self.history[indices_or_msg_ids]
         elif isinstance(indices_or_msg_ids, (list,tuple,set)):
             indices_or_msg_ids = list(indices_or_msg_ids)
             for i,index in enumerate(indices_or_msg_ids):
                 if isinstance(index, int):
                     indices_or_msg_ids[i] = self.history[index]
         return self.client.get_result(indices_or_msg_ids, block=block, owner=owner)
 
     #-------------------------------------------------------------------
     # Map
     #-------------------------------------------------------------------
 
     @sync_results
     def map(self, f, *sequences, **kwargs):
         """override in subclasses"""
         raise NotImplementedError
 
     def map_async(self, f, *sequences, **kwargs):
         """Parallel version of builtin :func:`python:map`, using this view's engines.
 
         This is equivalent to ``map(...block=False)``.
 
         See `self.map` for details.
         """
         if 'block' in kwargs:
             raise TypeError("map_async doesn't take a `block` keyword argument.")
         kwargs['block'] = False
         return self.map(f,*sequences,**kwargs)
 
     def map_sync(self, f, *sequences, **kwargs):
         """Parallel version of builtin :func:`python:map`, using this view's engines.
 
         This is equivalent to ``map(...block=True)``.
 
         See `self.map` for details.
         """
         if 'block' in kwargs:
             raise TypeError("map_sync doesn't take a `block` keyword argument.")
         kwargs['block'] = True
         return self.map(f,*sequences,**kwargs)
 
     def imap(self, f, *sequences, **kwargs):
         """Parallel version of :func:`itertools.imap`.
 
         See `self.map` for details.
 
         """
 
         return iter(self.map_async(f,*sequences, **kwargs))
 
     #-------------------------------------------------------------------
     # Decorators
     #-------------------------------------------------------------------
 
     def remote(self, block=None, **flags):
         """Decorator for making a RemoteFunction"""
         block = self.block if block is None else block
         return remote(self, block=block, **flags)
 
     def parallel(self, dist='b', block=None, **flags):
         """Decorator for making a ParallelFunction"""
         block = self.block if block is None else block
         return parallel(self, dist=dist, block=block, **flags)
 
 @skip_doctest
 class DirectView(View):
     """Direct Multiplexer View of one or more engines.
 
     These are created via indexed access to a client:
 
     >>> dv_1 = client[1]
     >>> dv_all = client[:]
     >>> dv_even = client[::2]
     >>> dv_some = client[1:3]
 
     This object provides dictionary access to engine namespaces:
 
     # push a=5:
     >>> dv['a'] = 5
     # pull 'foo':
     >>> dv['foo']
 
     """
 
     def __init__(self, client=None, socket=None, targets=None):
         super(DirectView, self).__init__(client=client, socket=socket, targets=targets)
 
     @property
     def importer(self):
         """sync_imports(local=True) as a property.
 
         See sync_imports for details.
 
         """
         return self.sync_imports(True)
 
     @contextmanager
     def sync_imports(self, local=True, quiet=False):
         """Context Manager for performing simultaneous local and remote imports.
 
         'import x as y' will *not* work.  The 'as y' part will simply be ignored.
         
         If `local=True`, then the package will also be imported locally.
         
         If `quiet=True`, no output will be produced when attempting remote 
         imports. 
         
         Note that remote-only (`local=False`) imports have not been implemented.
         
         >>> with view.sync_imports():
         ...    from numpy import recarray
         importing recarray from numpy on engine(s)
 
         """
         from IPython.utils.py3compat import builtin_mod
         local_import = builtin_mod.__import__
         modules = set()
         results = []
         @util.interactive
         def remote_import(name, fromlist, level):
             """the function to be passed to apply, that actually performs the import
             on the engine, and loads up the user namespace.
             """
             import sys
             user_ns = globals()
             mod = __import__(name, fromlist=fromlist, level=level)
             if fromlist:
                 for key in fromlist:
                     user_ns[key] = getattr(mod, key)
             else:
                 user_ns[name] = sys.modules[name]
 
         def view_import(name, globals={}, locals={}, fromlist=[], level=0):
             """the drop-in replacement for __import__, that optionally imports
             locally as well.
             """
             # don't override nested imports
             save_import = builtin_mod.__import__
             builtin_mod.__import__ = local_import
 
             if imp.lock_held():
                 # this is a side-effect import, don't do it remotely, or even
                 # ignore the local effects
                 return local_import(name, globals, locals, fromlist, level)
 
             imp.acquire_lock()
             if local:
                 mod = local_import(name, globals, locals, fromlist, level)
             else:
                 raise NotImplementedError("remote-only imports not yet implemented")
             imp.release_lock()
 
             key = name+':'+','.join(fromlist or [])
             if level <= 0 and key not in modules:
                 modules.add(key)
                 if not quiet:
                     if fromlist:
                         print("importing %s from %s on engine(s)"%(','.join(fromlist), name))
                     else:
                         print("importing %s on engine(s)"%name)
                 results.append(self.apply_async(remote_import, name, fromlist, level))
             # restore override
             builtin_mod.__import__ = save_import
 
             return mod
 
         # override __import__
         builtin_mod.__import__ = view_import
         try:
             # enter the block
             yield
         except ImportError:
             if local:
                 raise
             else:
                 # ignore import errors if not doing local imports
                 pass
         finally:
             # always restore __import__
             builtin_mod.__import__ = local_import
 
         for r in results:
             # raise possible remote ImportErrors here
             r.get()
     
     def use_dill(self):
         """Expand serialization support with dill
         
         adds support for closures, etc.
         
-        This calls IPython.utils.pickleutil.use_dill() here and on each engine.
+        This calls ipython_kernel.pickleutil.use_dill() here and on each engine.
         """
         pickleutil.use_dill()
         return self.apply(pickleutil.use_dill)
 
     def use_cloudpickle(self):
         """Expand serialization support with cloudpickle.
         """
         pickleutil.use_cloudpickle()
         return self.apply(pickleutil.use_cloudpickle)
 
 
     @sync_results
     @save_ids
     def _really_apply(self, f, args=None, kwargs=None, targets=None, block=None, track=None):
         """calls f(*args, **kwargs) on remote engines, returning the result.
 
         This method sets all of `apply`'s flags via this View's attributes.
 
         Parameters
         ----------
 
         f : callable
 
         args : list [default: empty]
 
         kwargs : dict [default: empty]
 
         targets : target list [default: self.targets]
             where to run
         block : bool [default: self.block]
             whether to block
         track : bool [default: self.track]
             whether to ask zmq to track the message, for safe non-copying sends
 
         Returns
         -------
 
         if self.block is False:
             returns AsyncResult
         else:
             returns actual result of f(*args, **kwargs) on the engine(s)
             This will be a list of self.targets is also a list (even length 1), or
             the single result if self.targets is an integer engine id
         """
         args = [] if args is None else args
         kwargs = {} if kwargs is None else kwargs
         block = self.block if block is None else block
         track = self.track if track is None else track
         targets = self.targets if targets is None else targets
         
         _idents, _targets = self.client._build_targets(targets)
         msg_ids = []
         trackers = []
         for ident in _idents:
             msg = self.client.send_apply_request(self._socket, f, args, kwargs, track=track,
                                     ident=ident)
             if track:
                 trackers.append(msg['tracker'])
             msg_ids.append(msg['header']['msg_id'])
         if isinstance(targets, int):
             msg_ids = msg_ids[0]
         tracker = None if track is False else zmq.MessageTracker(*trackers)
         ar = AsyncResult(self.client, msg_ids, fname=getname(f), targets=_targets,
             tracker=tracker, owner=True,
         )
         if block:
             try:
                 return ar.get()
             except KeyboardInterrupt:
                 pass
         return ar
 
 
     @sync_results
     def map(self, f, *sequences, **kwargs):
         """``view.map(f, *sequences, block=self.block)`` => list|AsyncMapResult
 
         Parallel version of builtin `map`, using this View's `targets`.
 
         There will be one task per target, so work will be chunked
         if the sequences are longer than `targets`.
 
         Results can be iterated as they are ready, but will become available in chunks.
 
         Parameters
         ----------
 
         f : callable
             function to be mapped
         *sequences: one or more sequences of matching length
             the sequences to be distributed and passed to `f`
         block : bool
             whether to wait for the result or not [default self.block]
 
         Returns
         -------
 
 
         If block=False
           An :class:`~ipython_parallel.client.asyncresult.AsyncMapResult` instance.
           An object like AsyncResult, but which reassembles the sequence of results
           into a single list. AsyncMapResults can be iterated through before all
           results are complete.
         else
             A list, the result of ``map(f,*sequences)``
         """
 
         block = kwargs.pop('block', self.block)
         for k in kwargs.keys():
             if k not in ['block', 'track']:
                 raise TypeError("invalid keyword arg, %r"%k)
 
         assert len(sequences) > 0, "must have some sequences to map onto!"
         pf = ParallelFunction(self, f, block=block, **kwargs)
         return pf.map(*sequences)
 
     @sync_results
     @save_ids
     def execute(self, code, silent=True, targets=None, block=None):
         """Executes `code` on `targets` in blocking or nonblocking manner.
 
         ``execute`` is always `bound` (affects engine namespace)
 
         Parameters
         ----------
 
         code : str
                 the code string to be executed
         block : bool
                 whether or not to wait until done to return
                 default: self.block
         """
         block = self.block if block is None else block
         targets = self.targets if targets is None else targets
 
         _idents, _targets = self.client._build_targets(targets)
         msg_ids = []
         trackers = []
         for ident in _idents:
             msg = self.client.send_execute_request(self._socket, code, silent=silent, ident=ident)
             msg_ids.append(msg['header']['msg_id'])
         if isinstance(targets, int):
             msg_ids = msg_ids[0]
         ar = AsyncResult(self.client, msg_ids, fname='execute', targets=_targets, owner=True)
         if block:
             try:
                 ar.get()
             except KeyboardInterrupt:
                 pass
         return ar
 
     def run(self, filename, targets=None, block=None):
         """Execute contents of `filename` on my engine(s).
 
         This simply reads the contents of the file and calls `execute`.
 
         Parameters
         ----------
 
         filename : str
                 The path to the file
         targets : int/str/list of ints/strs
                 the engines on which to execute
                 default : all
         block : bool
                 whether or not to wait until done
                 default: self.block
 
         """
         with open(filename, 'r') as f:
             # add newline in case of trailing indented whitespace
             # which will cause SyntaxError
             code = f.read()+'\n'
         return self.execute(code, block=block, targets=targets)
 
     def update(self, ns):
         """update remote namespace with dict `ns`
 
         See `push` for details.
         """
         return self.push(ns, block=self.block, track=self.track)
 
     def push(self, ns, targets=None, block=None, track=None):
         """update remote namespace with dict `ns`
 
         Parameters
         ----------
 
         ns : dict
             dict of keys with which to update engine namespace(s)
         block : bool [default : self.block]
             whether to wait to be notified of engine receipt
 
         """
 
         block = block if block is not None else self.block
         track = track if track is not None else self.track
         targets = targets if targets is not None else self.targets
         # applier = self.apply_sync if block else self.apply_async
         if not isinstance(ns, dict):
             raise TypeError("Must be a dict, not %s"%type(ns))
         return self._really_apply(util._push, kwargs=ns, block=block, track=track, targets=targets)
 
     def get(self, key_s):
         """get object(s) by `key_s` from remote namespace
 
         see `pull` for details.
         """
         # block = block if block is not None else self.block
         return self.pull(key_s, block=True)
 
     def pull(self, names, targets=None, block=None):
         """get object(s) by `name` from remote namespace
 
         will return one object if it is a key.
         can also take a list of keys, in which case it will return a list of objects.
         """
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         applier = self.apply_sync if block else self.apply_async
         if isinstance(names, string_types):
             pass
         elif isinstance(names, (list,tuple,set)):
             for key in names:
                 if not isinstance(key, string_types):
                     raise TypeError("keys must be str, not type %r"%type(key))
         else:
             raise TypeError("names must be strs, not %r"%names)
         return self._really_apply(util._pull, (names,), block=block, targets=targets)
 
     def scatter(self, key, seq, dist='b', flatten=False, targets=None, block=None, track=None):
         """
         Partition a Python sequence and send the partitions to a set of engines.
         """
         block = block if block is not None else self.block
         track = track if track is not None else self.track
         targets = targets if targets is not None else self.targets
         
         # construct integer ID list:
         targets = self.client._build_targets(targets)[1]
 
         mapObject = Map.dists[dist]()
         nparts = len(targets)
         msg_ids = []
         trackers = []
         for index, engineid in enumerate(targets):
             partition = mapObject.getPartition(seq, index, nparts)
             if flatten and len(partition) == 1:
                 ns = {key: partition[0]}
             else:
                 ns = {key: partition}
             r = self.push(ns, block=False, track=track, targets=engineid)
             msg_ids.extend(r.msg_ids)
             if track:
                 trackers.append(r._tracker)
 
         if track:
             tracker = zmq.MessageTracker(*trackers)
         else:
             tracker = None
 
         r = AsyncResult(self.client, msg_ids, fname='scatter', targets=targets,
             tracker=tracker, owner=True,
         )
         if block:
             r.wait()
         else:
             return r
 
     @sync_results
     @save_ids
     def gather(self, key, dist='b', targets=None, block=None):
         """
         Gather a partitioned sequence on a set of engines as a single local seq.
         """
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         mapObject = Map.dists[dist]()
         msg_ids = []
 
         # construct integer ID list:
         targets = self.client._build_targets(targets)[1]
         
         for index, engineid in enumerate(targets):
             msg_ids.extend(self.pull(key, block=False, targets=engineid).msg_ids)
 
         r = AsyncMapResult(self.client, msg_ids, mapObject, fname='gather')
 
         if block:
             try:
                 return r.get()
             except KeyboardInterrupt:
                 pass
         return r
 
     def __getitem__(self, key):
         return self.get(key)
 
     def __setitem__(self,key, value):
         self.update({key:value})
 
     def clear(self, targets=None, block=None):
         """Clear the remote namespaces on my engines."""
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         return self.client.clear(targets=targets, block=block)
 
     #----------------------------------------
     # activate for %px, %autopx, etc. magics
     #----------------------------------------
 
     def activate(self, suffix=''):
         """Activate IPython magics associated with this View
         
         Defines the magics `%px, %autopx, %pxresult, %%px, %pxconfig`
         
         Parameters
         ----------
         
         suffix: str [default: '']
             The suffix, if any, for the magics.  This allows you to have
             multiple views associated with parallel magics at the same time.
             
             e.g. ``rc[::2].activate(suffix='_even')`` will give you
             the magics ``%px_even``, ``%pxresult_even``, etc. for running magics 
             on the even engines.
         """
         
         from IPython.parallel.client.magics import ParallelMagics
         
         try:
             # This is injected into __builtins__.
             ip = get_ipython()
         except NameError:
             print("The IPython parallel magics (%px, etc.) only work within IPython.")
             return
         
         M = ParallelMagics(ip, self, suffix)
         ip.magics_manager.register(M)
 
 
 @skip_doctest
 class LoadBalancedView(View):
     """An load-balancing View that only executes via the Task scheduler.
 
     Load-balanced views can be created with the client's `view` method:
 
     >>> v = client.load_balanced_view()
 
     or targets can be specified, to restrict the potential destinations:
 
     >>> v = client.load_balanced_view([1,3])
 
     which would restrict loadbalancing to between engines 1 and 3.
 
     """
 
     follow=Any()
     after=Any()
     timeout=CFloat()
     retries = Integer(0)
 
     _task_scheme = Any()
     _flag_names = List(['targets', 'block', 'track', 'follow', 'after', 'timeout', 'retries'])
 
     def __init__(self, client=None, socket=None, **flags):
         super(LoadBalancedView, self).__init__(client=client, socket=socket, **flags)
         self._task_scheme=client._task_scheme
 
     def _validate_dependency(self, dep):
         """validate a dependency.
 
         For use in `set_flags`.
         """
         if dep is None or isinstance(dep, string_types + (AsyncResult, Dependency)):
             return True
         elif isinstance(dep, (list,set, tuple)):
             for d in dep:
                 if not isinstance(d, string_types + (AsyncResult,)):
                     return False
         elif isinstance(dep, dict):
             if set(dep.keys()) != set(Dependency().as_dict().keys()):
                 return False
             if not isinstance(dep['msg_ids'], list):
                 return False
             for d in dep['msg_ids']:
                 if not isinstance(d, string_types):
                     return False
         else:
             return False
 
         return True
 
     def _render_dependency(self, dep):
         """helper for building jsonable dependencies from various input forms."""
         if isinstance(dep, Dependency):
             return dep.as_dict()
         elif isinstance(dep, AsyncResult):
             return dep.msg_ids
         elif dep is None:
             return []
         else:
             # pass to Dependency constructor
             return list(Dependency(dep))
 
     def set_flags(self, **kwargs):
         """set my attribute flags by keyword.
 
         A View is a wrapper for the Client's apply method, but with attributes
         that specify keyword arguments, those attributes can be set by keyword
         argument with this method.
 
         Parameters
         ----------
 
         block : bool
             whether to wait for results
         track : bool
             whether to create a MessageTracker to allow the user to
             safely edit after arrays and buffers during non-copying
             sends.
 
         after : Dependency or collection of msg_ids
             Only for load-balanced execution (targets=None)
             Specify a list of msg_ids as a time-based dependency.
             This job will only be run *after* the dependencies
             have been met.
 
         follow : Dependency or collection of msg_ids
             Only for load-balanced execution (targets=None)
             Specify a list of msg_ids as a location-based dependency.
             This job will only be run on an engine where this dependency
             is met.
 
         timeout : float/int or None
             Only for load-balanced execution (targets=None)
             Specify an amount of time (in seconds) for the scheduler to
             wait for dependencies to be met before failing with a
             DependencyTimeout.
 
         retries : int
             Number of times a task will be retried on failure.
         """
 
         super(LoadBalancedView, self).set_flags(**kwargs)
         for name in ('follow', 'after'):
             if name in kwargs:
                 value = kwargs[name]
                 if self._validate_dependency(value):
                     setattr(self, name, value)
                 else:
                     raise ValueError("Invalid dependency: %r"%value)
         if 'timeout' in kwargs:
             t = kwargs['timeout']
             if not isinstance(t, (int, float, type(None))):
                 if (not PY3) and (not isinstance(t, long)):
                     raise TypeError("Invalid type for timeout: %r"%type(t))
             if t is not None:
                 if t < 0:
                     raise ValueError("Invalid timeout: %s"%t)
             self.timeout = t
 
     @sync_results
     @save_ids
     def _really_apply(self, f, args=None, kwargs=None, block=None, track=None,
                                         after=None, follow=None, timeout=None,
                                         targets=None, retries=None):
         """calls f(*args, **kwargs) on a remote engine, returning the result.
 
         This method temporarily sets all of `apply`'s flags for a single call.
 
         Parameters
         ----------
 
         f : callable
 
         args : list [default: empty]
 
         kwargs : dict [default: empty]
 
         block : bool [default: self.block]
             whether to block
         track : bool [default: self.track]
             whether to ask zmq to track the message, for safe non-copying sends
 
         !!!!!! TODO: THE REST HERE  !!!!
 
         Returns
         -------
 
         if self.block is False:
             returns AsyncResult
         else:
             returns actual result of f(*args, **kwargs) on the engine(s)
             This will be a list of self.targets is also a list (even length 1), or
             the single result if self.targets is an integer engine id
         """
 
         # validate whether we can run
         if self._socket.closed:
             msg = "Task farming is disabled"
             if self._task_scheme == 'pure':
                 msg += " because the pure ZMQ scheduler cannot handle"
                 msg += " disappearing engines."
             raise RuntimeError(msg)
 
         if self._task_scheme == 'pure':
             # pure zmq scheme doesn't support extra features
             msg = "Pure ZMQ scheduler doesn't support the following flags:"
             "follow, after, retries, targets, timeout"
             if (follow or after or retries or targets or timeout):
                 # hard fail on Scheduler flags
                 raise RuntimeError(msg)
             if isinstance(f, dependent):
                 # soft warn on functional dependencies
                 warnings.warn(msg, RuntimeWarning)
 
         # build args
         args = [] if args is None else args
         kwargs = {} if kwargs is None else kwargs
         block = self.block if block is None else block
         track = self.track if track is None else track
         after = self.after if after is None else after
         retries = self.retries if retries is None else retries
         follow = self.follow if follow is None else follow
         timeout = self.timeout if timeout is None else timeout
         targets = self.targets if targets is None else targets
 
         if not isinstance(retries, int):
             raise TypeError('retries must be int, not %r'%type(retries))
 
         if targets is None:
             idents = []
         else:
             idents = self.client._build_targets(targets)[0]
             # ensure *not* bytes
             idents = [ ident.decode() for ident in idents ]
 
         after = self._render_dependency(after)
         follow = self._render_dependency(follow)
         metadata = dict(after=after, follow=follow, timeout=timeout, targets=idents, retries=retries)
 
         msg = self.client.send_apply_request(self._socket, f, args, kwargs, track=track,
                                 metadata=metadata)
         tracker = None if track is False else msg['tracker']
 
         ar = AsyncResult(self.client, msg['header']['msg_id'], fname=getname(f),
             targets=None, tracker=tracker, owner=True,
         )
         if block:
             try:
                 return ar.get()
             except KeyboardInterrupt:
                 pass
         return ar
 
     @sync_results
     @save_ids
     def map(self, f, *sequences, **kwargs):
         """``view.map(f, *sequences, block=self.block, chunksize=1, ordered=True)`` => list|AsyncMapResult
 
         Parallel version of builtin `map`, load-balanced by this View.
 
         `block`, and `chunksize` can be specified by keyword only.
 
         Each `chunksize` elements will be a separate task, and will be
         load-balanced. This lets individual elements be available for iteration
         as soon as they arrive.
 
         Parameters
         ----------
 
         f : callable
             function to be mapped
         *sequences: one or more sequences of matching length
             the sequences to be distributed and passed to `f`
         block : bool [default self.block]
             whether to wait for the result or not
         track : bool
             whether to create a MessageTracker to allow the user to
             safely edit after arrays and buffers during non-copying
             sends.
         chunksize : int [default 1]
             how many elements should be in each task.
         ordered : bool [default True]
             Whether the results should be gathered as they arrive, or enforce
             the order of submission.
             
             Only applies when iterating through AsyncMapResult as results arrive.
             Has no effect when block=True.
 
         Returns
         -------
 
         if block=False
           An :class:`~ipython_parallel.client.asyncresult.AsyncMapResult` instance.
           An object like AsyncResult, but which reassembles the sequence of results
           into a single list. AsyncMapResults can be iterated through before all
           results are complete.
         else
             A list, the result of ``map(f,*sequences)``
         """
 
         # default
         block = kwargs.get('block', self.block)
         chunksize = kwargs.get('chunksize', 1)
         ordered = kwargs.get('ordered', True)
 
         keyset = set(kwargs.keys())
         extra_keys = keyset.difference_update(set(['block', 'chunksize']))
         if extra_keys:
             raise TypeError("Invalid kwargs: %s"%list(extra_keys))
 
         assert len(sequences) > 0, "must have some sequences to map onto!"
 
         pf = ParallelFunction(self, f, block=block, chunksize=chunksize, ordered=ordered)
         return pf.map(*sequences)
 
 __all__ = ['LoadBalancedView', 'DirectView']

ipython_parallel/controller/dependency.py

@@ -1,229 +1,229 b''
 """Dependency utilities
 
 Authors:
 
 * Min RK
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2013  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 from types import ModuleType
 
 from ipython_parallel.client.asyncresult import AsyncResult
 from ipython_parallel.error import UnmetDependency
 from ipython_parallel.util import interactive
 from IPython.utils import py3compat
 from IPython.utils.py3compat import string_types
-from IPython.utils.pickleutil import can, uncan
+from ipython_kernel.pickleutil import can, uncan
 
 class depend(object):
     """Dependency decorator, for use with tasks.
     
     `@depend` lets you define a function for engine dependencies
     just like you use `apply` for tasks.
     
     
     Examples
     --------
     ::
     
         @depend(df, a,b, c=5)
         def f(m,n,p)
         
         view.apply(f, 1,2,3)
     
     will call df(a,b,c=5) on the engine, and if it returns False or
     raises an UnmetDependency error, then the task will not be run
     and another engine will be tried.
     """
     def __init__(self, _wrapped_f, *args, **kwargs):
         self.f = _wrapped_f
         self.args = args
         self.kwargs = kwargs
     
     def __call__(self, f):
         return dependent(f, self.f, *self.args, **self.kwargs)
 
 class dependent(object):
     """A function that depends on another function.
     This is an object to prevent the closure used
     in traditional decorators, which are not picklable.
     """
 
     def __init__(self, _wrapped_f, _wrapped_df, *dargs, **dkwargs):
         self.f = _wrapped_f
         name = getattr(_wrapped_f, '__name__', 'f')
         if py3compat.PY3:
             self.__name__ = name
         else:
             self.func_name = name
         self.df = _wrapped_df
         self.dargs = dargs
         self.dkwargs = dkwargs
     
     def check_dependency(self):
         if self.df(*self.dargs, **self.dkwargs) is False:
             raise UnmetDependency()
     
     def __call__(self, *args, **kwargs):
         return self.f(*args, **kwargs)
     
     if not py3compat.PY3:
         @property
         def __name__(self):
             return self.func_name
 
 @interactive
 def _require(*modules, **mapping):
     """Helper for @require decorator."""
     from ipython_parallel.error import UnmetDependency
-    from IPython.utils.pickleutil import uncan
+    from ipython_kernel.pickleutil import uncan
     user_ns = globals()
     for name in modules:
         try:
             exec('import %s' % name, user_ns)
         except ImportError:
             raise UnmetDependency(name)
             
     for name, cobj in mapping.items():
         user_ns[name] = uncan(cobj, user_ns)
     return True
 
 def require(*objects, **mapping):
     """Simple decorator for requiring local objects and modules to be available
     when the decorated function is called on the engine.
     
     Modules specified by name or passed directly will be imported
     prior to calling the decorated function.
     
     Objects other than modules will be pushed as a part of the task.
     Functions can be passed positionally,
     and will be pushed to the engine with their __name__.
     Other objects can be passed by keyword arg.
     
     Examples::
     
         In [1]: @require('numpy')
            ...: def norm(a):
            ...:     return numpy.linalg.norm(a,2)
 
         In [2]: foo = lambda x: x*x
         In [3]: @require(foo)
            ...: def bar(a):
            ...:     return foo(1-a)
     """
     names = []
     for obj in objects:
         if isinstance(obj, ModuleType):
             obj = obj.__name__
         
         if isinstance(obj, string_types):
             names.append(obj)
         elif hasattr(obj, '__name__'):
             mapping[obj.__name__] = obj
         else:
             raise TypeError("Objects other than modules and functions "
                 "must be passed by kwarg, but got: %s" % type(obj)
             )
     
     for name, obj in mapping.items():
         mapping[name] = can(obj)
     return depend(_require, *names, **mapping)
 
 class Dependency(set):
     """An object for representing a set of msg_id dependencies.
     
     Subclassed from set().
     
     Parameters
     ----------
     dependencies: list/set of msg_ids or AsyncResult objects or output of Dependency.as_dict()
         The msg_ids to depend on
     all : bool [default True]
         Whether the dependency should be considered met when *all* depending tasks have completed
         or only when *any* have been completed.
     success : bool [default True]
         Whether to consider successes as fulfilling dependencies.
     failure : bool [default False]
         Whether to consider failures as fulfilling dependencies.
     
     If `all=success=True` and `failure=False`, then the task will fail with an ImpossibleDependency
         as soon as the first depended-upon task fails.
     """
     
     all=True
     success=True
     failure=True
     
     def __init__(self, dependencies=[], all=True, success=True, failure=False):
         if isinstance(dependencies, dict):
             # load from dict
             all = dependencies.get('all', True)
             success = dependencies.get('success', success)
             failure = dependencies.get('failure', failure)
             dependencies = dependencies.get('dependencies', [])
         ids = []
         
         # extract ids from various sources:
         if isinstance(dependencies, string_types + (AsyncResult,)):
             dependencies = [dependencies]
         for d in dependencies:
             if isinstance(d, string_types):
                 ids.append(d)
             elif isinstance(d, AsyncResult):
                 ids.extend(d.msg_ids)
             else:
                 raise TypeError("invalid dependency type: %r"%type(d))
         
         set.__init__(self, ids)
         self.all = all
         if not (success or failure):
             raise ValueError("Must depend on at least one of successes or failures!")
         self.success=success
         self.failure = failure
     
     def check(self, completed, failed=None):
         """check whether our dependencies have been met."""
         if len(self) == 0:
             return True
         against = set()
         if self.success:
             against = completed
         if failed is not None and self.failure:
             against = against.union(failed)
         if self.all:
             return self.issubset(against)
         else:
             return not self.isdisjoint(against)
     
     def unreachable(self, completed, failed=None):
         """return whether this dependency has become impossible."""
         if len(self) == 0:
             return False
         against = set()
         if not self.success:
             against = completed
         if failed is not None and not self.failure:
             against = against.union(failed)
         if self.all:
             return not self.isdisjoint(against)
         else:
             return self.issubset(against)
         
     
     def as_dict(self):
         """Represent this dependency as a dict. For json compatibility."""
         return dict(
             dependencies=list(self),
             all=self.all,
             success=self.success,
             failure=self.failure
         )
 
 
 __all__ = ['depend', 'require', 'dependent', 'Dependency']
 

ipython_parallel/tests/test_dependency.py

@@ -1,136 +1,136 b''
 """Tests for dependency.py
 
 Authors:
 
 * Min RK
 """
 
 __docformat__ = "restructuredtext en"
 
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
 # import
 import os
 
-from IPython.utils.pickleutil import can, uncan
+from ipython_kernel.pickleutil import can, uncan
 
 import ipython_parallel as pmod
 from ipython_parallel.util import interactive
 
 from ipython_parallel.tests import add_engines
 from .clienttest import ClusterTestCase
 
 def setup():
     add_engines(1, total=True)
 
 @pmod.require('time')
 def wait(n):
     time.sleep(n)
     return n
 
 @pmod.interactive
 def func(x):
     return x*x
 
 mixed = list(map(str, range(10)))
 completed = list(map(str, range(0,10,2)))
 failed = list(map(str, range(1,10,2)))
 
 class DependencyTest(ClusterTestCase):
     
     def setUp(self):
         ClusterTestCase.setUp(self)
         self.user_ns = {'__builtins__' : __builtins__}
         self.view = self.client.load_balanced_view()
         self.dview = self.client[-1]
         self.succeeded = set(map(str, range(0,25,2)))
         self.failed = set(map(str, range(1,25,2)))
     
     def assertMet(self, dep):
         self.assertTrue(dep.check(self.succeeded, self.failed), "Dependency should be met")
         
     def assertUnmet(self, dep):
         self.assertFalse(dep.check(self.succeeded, self.failed), "Dependency should not be met")
         
     def assertUnreachable(self, dep):
         self.assertTrue(dep.unreachable(self.succeeded, self.failed), "Dependency should be unreachable")
     
     def assertReachable(self, dep):
         self.assertFalse(dep.unreachable(self.succeeded, self.failed), "Dependency should be reachable")
     
     def cancan(self, f):
         """decorator to pass through canning into self.user_ns"""
         return uncan(can(f), self.user_ns)
     
     def test_require_imports(self):
         """test that @require imports names"""
         @self.cancan
         @pmod.require('base64')
         @interactive
         def encode(arg):
             return base64.b64encode(arg)
         # must pass through canning to properly connect namespaces
         self.assertEqual(encode(b'foo'), b'Zm9v')
     
     def test_success_only(self):
         dep = pmod.Dependency(mixed, success=True, failure=False)
         self.assertUnmet(dep)
         self.assertUnreachable(dep)
         dep.all=False
         self.assertMet(dep)
         self.assertReachable(dep)
         dep = pmod.Dependency(completed, success=True, failure=False)
         self.assertMet(dep)
         self.assertReachable(dep)
         dep.all=False
         self.assertMet(dep)
         self.assertReachable(dep)
 
     def test_failure_only(self):
         dep = pmod.Dependency(mixed, success=False, failure=True)
         self.assertUnmet(dep)
         self.assertUnreachable(dep)
         dep.all=False
         self.assertMet(dep)
         self.assertReachable(dep)
         dep = pmod.Dependency(completed, success=False, failure=True)
         self.assertUnmet(dep)
         self.assertUnreachable(dep)
         dep.all=False
         self.assertUnmet(dep)
         self.assertUnreachable(dep)
     
     def test_require_function(self):
         
         @pmod.interactive
         def bar(a):
             return func(a)
 
         @pmod.require(func)
         @pmod.interactive
         def bar2(a):
             return func(a)
         
         self.client[:].clear()
         self.assertRaisesRemote(NameError, self.view.apply_sync, bar, 5)
         ar = self.view.apply_async(bar2, 5)
         self.assertEqual(ar.get(5), func(5))
 
     def test_require_object(self):
         
         @pmod.require(foo=func)
         @pmod.interactive
         def bar(a):
             return foo(a)
 
         ar = self.view.apply_async(bar, 5)
         self.assertEqual(ar.get(5), func(5))

jupyter_client/session.py

@@ -1,884 +1,889 b''
 """Session object for building, serializing, sending, and receiving messages in
 IPython. The Session object supports serialization, HMAC signatures, and
 metadata on messages.
 
 Also defined here are utilities for working with Sessions:
 * A SessionFactory to be used as a base class for configurables that work with
 Sessions.
 * A Message object for convenience that allows attribute-access to the msg dict.
 """
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import hashlib
 import hmac
 import logging
 import os
 import pprint
 import random
 import uuid
 import warnings
 from datetime import datetime
 
 try:
     import cPickle
     pickle = cPickle
 except:
     cPickle = None
     import pickle
 
 try:
+    # py3
+    PICKLE_PROTOCOL = pickle.DEFAULT_PROTOCOL
+except AttributeError:
+    PICKLE_PROTOCOL = pickle.HIGHEST_PROTOCOL
+
+try:
     # We are using compare_digest to limit the surface of timing attacks
     from hmac import compare_digest
 except ImportError:
     # Python < 2.7.7: When digests don't match no feedback is provided,
     # limiting the surface of attack
     def compare_digest(a,b): return a == b
 
 import zmq
 from zmq.utils import jsonapi
 from zmq.eventloop.ioloop import IOLoop
 from zmq.eventloop.zmqstream import ZMQStream
 
 from IPython.core.release import kernel_protocol_version
 from IPython.config.configurable import Configurable, LoggingConfigurable
 from IPython.utils import io
 from IPython.utils.importstring import import_item
 from IPython.utils.jsonutil import extract_dates, squash_dates, date_default
 from IPython.utils.py3compat import (str_to_bytes, str_to_unicode, unicode_type,
                                      iteritems)
 from IPython.utils.traitlets import (CBytes, Unicode, Bool, Any, Instance, Set,
                                         DottedObjectName, CUnicode, Dict, Integer,
                                         TraitError,
 )
-from IPython.utils.pickleutil import PICKLE_PROTOCOL
 from jupyter_client.adapter import adapt
 
 #-----------------------------------------------------------------------------
 # utility functions
 #-----------------------------------------------------------------------------
 
 def squash_unicode(obj):
     """coerce unicode back to bytestrings."""
     if isinstance(obj,dict):
         for key in obj.keys():
             obj[key] = squash_unicode(obj[key])
             if isinstance(key, unicode_type):
                 obj[squash_unicode(key)] = obj.pop(key)
     elif isinstance(obj, list):
         for i,v in enumerate(obj):
             obj[i] = squash_unicode(v)
     elif isinstance(obj, unicode_type):
         obj = obj.encode('utf8')
     return obj
 
 #-----------------------------------------------------------------------------
 # globals and defaults
 #-----------------------------------------------------------------------------
 
 # default values for the thresholds:
 MAX_ITEMS = 64
 MAX_BYTES = 1024
 
 # ISO8601-ify datetime objects
 # allow unicode
 # disallow nan, because it's not actually valid JSON
 json_packer = lambda obj: jsonapi.dumps(obj, default=date_default,
     ensure_ascii=False, allow_nan=False,
 )
 json_unpacker = lambda s: jsonapi.loads(s)
 
 pickle_packer = lambda o: pickle.dumps(squash_dates(o), PICKLE_PROTOCOL)
 pickle_unpacker = pickle.loads
 
 default_packer = json_packer
 default_unpacker = json_unpacker
 
 DELIM = b"<IDS|MSG>"
 # singleton dummy tracker, which will always report as done
 DONE = zmq.MessageTracker()
 
 #-----------------------------------------------------------------------------
 # Mixin tools for apps that use Sessions
 #-----------------------------------------------------------------------------
 
 session_aliases = dict(
     ident = 'Session.session',
     user = 'Session.username',
     keyfile = 'Session.keyfile',
 )
 
 session_flags  = {
     'secure' : ({'Session' : { 'key' : str_to_bytes(str(uuid.uuid4())),
                             'keyfile' : '' }},
         """Use HMAC digests for authentication of messages.
         Setting this flag will generate a new UUID to use as the HMAC key.
         """),
     'no-secure' : ({'Session' : { 'key' : b'', 'keyfile' : '' }},
         """Don't authenticate messages."""),
 }
 
 def default_secure(cfg):
     """Set the default behavior for a config environment to be secure.
 
     If Session.key/keyfile have not been set, set Session.key to
     a new random UUID.
     """
     warnings.warn("default_secure is deprecated", DeprecationWarning)
     if 'Session' in cfg:
         if 'key' in cfg.Session or 'keyfile' in cfg.Session:
             return
     # key/keyfile not specified, generate new UUID:
     cfg.Session.key = str_to_bytes(str(uuid.uuid4()))
 
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
 class SessionFactory(LoggingConfigurable):
     """The Base class for configurables that have a Session, Context, logger,
     and IOLoop.
     """
 
     logname = Unicode('')
     def _logname_changed(self, name, old, new):
         self.log = logging.getLogger(new)
 
     # not configurable:
     context = Instance('zmq.Context')
     def _context_default(self):
         return zmq.Context.instance()
 
     session = Instance('jupyter_client.session.Session',
                        allow_none=True)
 
     loop = Instance('zmq.eventloop.ioloop.IOLoop')
     def _loop_default(self):
         return IOLoop.instance()
 
     def __init__(self, **kwargs):
         super(SessionFactory, self).__init__(**kwargs)
 
         if self.session is None:
             # construct the session
             self.session = Session(**kwargs)
 
 
 class Message(object):
     """A simple message object that maps dict keys to attributes.
 
     A Message can be created from a dict and a dict from a Message instance
     simply by calling dict(msg_obj)."""
 
     def __init__(self, msg_dict):
         dct = self.__dict__
         for k, v in iteritems(dict(msg_dict)):
             if isinstance(v, dict):
                 v = Message(v)
             dct[k] = v
 
     # Having this iterator lets dict(msg_obj) work out of the box.
     def __iter__(self):
         return iter(iteritems(self.__dict__))
 
     def __repr__(self):
         return repr(self.__dict__)
 
     def __str__(self):
         return pprint.pformat(self.__dict__)
 
     def __contains__(self, k):
         return k in self.__dict__
 
     def __getitem__(self, k):
         return self.__dict__[k]
 
 
 def msg_header(msg_id, msg_type, username, session):
     date = datetime.now()
     version = kernel_protocol_version
     return locals()
 
 def extract_header(msg_or_header):
     """Given a message or header, return the header."""
     if not msg_or_header:
         return {}
     try:
         # See if msg_or_header is the entire message.
         h = msg_or_header['header']
     except KeyError:
         try:
             # See if msg_or_header is just the header
             h = msg_or_header['msg_id']
         except KeyError:
             raise
         else:
             h = msg_or_header
     if not isinstance(h, dict):
         h = dict(h)
     return h
 
 class Session(Configurable):
     """Object for handling serialization and sending of messages.
 
     The Session object handles building messages and sending them
     with ZMQ sockets or ZMQStream objects.  Objects can communicate with each
     other over the network via Session objects, and only need to work with the
     dict-based IPython message spec. The Session will handle
     serialization/deserialization, security, and metadata.
 
     Sessions support configurable serialization via packer/unpacker traits,
     and signing with HMAC digests via the key/keyfile traits.
 
     Parameters
     ----------
 
     debug : bool
         whether to trigger extra debugging statements
     packer/unpacker : str : 'json', 'pickle' or import_string
         importstrings for methods to serialize message parts.  If just
         'json' or 'pickle', predefined JSON and pickle packers will be used.
         Otherwise, the entire importstring must be used.
 
         The functions must accept at least valid JSON input, and output *bytes*.
 
         For example, to use msgpack:
         packer = 'msgpack.packb', unpacker='msgpack.unpackb'
     pack/unpack : callables
         You can also set the pack/unpack callables for serialization directly.
     session : bytes
         the ID of this Session object.  The default is to generate a new UUID.
     username : unicode
         username added to message headers.  The default is to ask the OS.
     key : bytes
         The key used to initialize an HMAC signature.  If unset, messages
         will not be signed or checked.
     keyfile : filepath
         The file containing a key.  If this is set, `key` will be initialized
         to the contents of the file.
 
     """
 
     debug=Bool(False, config=True, help="""Debug output in the Session""")
 
     packer = DottedObjectName('json',config=True,
             help="""The name of the packer for serializing messages.
             Should be one of 'json', 'pickle', or an import name
             for a custom callable serializer.""")
     def _packer_changed(self, name, old, new):
         if new.lower() == 'json':
             self.pack = json_packer
             self.unpack = json_unpacker
             self.unpacker = new
         elif new.lower() == 'pickle':
             self.pack = pickle_packer
             self.unpack = pickle_unpacker
             self.unpacker = new
         else:
             self.pack = import_item(str(new))
 
     unpacker = DottedObjectName('json', config=True,
         help="""The name of the unpacker for unserializing messages.
         Only used with custom functions for `packer`.""")
     def _unpacker_changed(self, name, old, new):
         if new.lower() == 'json':
             self.pack = json_packer
             self.unpack = json_unpacker
             self.packer = new
         elif new.lower() == 'pickle':
             self.pack = pickle_packer
             self.unpack = pickle_unpacker
             self.packer = new
         else:
             self.unpack = import_item(str(new))
 
     session = CUnicode(u'', config=True,
         help="""The UUID identifying this session.""")
     def _session_default(self):
         u = unicode_type(uuid.uuid4())
         self.bsession = u.encode('ascii')
         return u
 
     def _session_changed(self, name, old, new):
         self.bsession = self.session.encode('ascii')
 
     # bsession is the session as bytes
     bsession = CBytes(b'')
 
     username = Unicode(str_to_unicode(os.environ.get('USER', 'username')),
         help="""Username for the Session. Default is your system username.""",
         config=True)
 
     metadata = Dict({}, config=True,
         help="""Metadata dictionary, which serves as the default top-level metadata dict for each message.""")
 
     # if 0, no adapting to do.
     adapt_version = Integer(0)
 
     # message signature related traits:
 
     key = CBytes(config=True,
         help="""execution key, for signing messages.""")
     def _key_default(self):
         return str_to_bytes(str(uuid.uuid4()))
 
     def _key_changed(self):
         self._new_auth()
 
     signature_scheme = Unicode('hmac-sha256', config=True,
         help="""The digest scheme used to construct the message signatures.
         Must have the form 'hmac-HASH'.""")
     def _signature_scheme_changed(self, name, old, new):
         if not new.startswith('hmac-'):
             raise TraitError("signature_scheme must start with 'hmac-', got %r" % new)
         hash_name = new.split('-', 1)[1]
         try:
             self.digest_mod = getattr(hashlib, hash_name)
         except AttributeError:
             raise TraitError("hashlib has no such attribute: %s" % hash_name)
         self._new_auth()
 
     digest_mod = Any()
     def _digest_mod_default(self):
         return hashlib.sha256
     
     auth = Instance(hmac.HMAC, allow_none=True)
     
     def _new_auth(self):
         if self.key:
             self.auth = hmac.HMAC(self.key, digestmod=self.digest_mod)
         else:
             self.auth = None
 
     digest_history = Set()
     digest_history_size = Integer(2**16, config=True,
         help="""The maximum number of digests to remember.
 
         The digest history will be culled when it exceeds this value.
         """
     )
 
     keyfile = Unicode('', config=True,
         help="""path to file containing execution key.""")
     def _keyfile_changed(self, name, old, new):
         with open(new, 'rb') as f:
             self.key = f.read().strip()
 
     # for protecting against sends from forks
     pid = Integer()
 
     # serialization traits:
 
     pack = Any(default_packer) # the actual packer function
     def _pack_changed(self, name, old, new):
         if not callable(new):
             raise TypeError("packer must be callable, not %s"%type(new))
 
     unpack = Any(default_unpacker) # the actual packer function
     def _unpack_changed(self, name, old, new):
         # unpacker is not checked - it is assumed to be
         if not callable(new):
             raise TypeError("unpacker must be callable, not %s"%type(new))
 
     # thresholds:
     copy_threshold = Integer(2**16, config=True,
         help="Threshold (in bytes) beyond which a buffer should be sent without copying.")
     buffer_threshold = Integer(MAX_BYTES, config=True,
         help="Threshold (in bytes) beyond which an object's buffer should be extracted to avoid pickling.")
     item_threshold = Integer(MAX_ITEMS, config=True,
         help="""The maximum number of items for a container to be introspected for custom serialization.
         Containers larger than this are pickled outright.
         """
     )
 
 
     def __init__(self, **kwargs):
         """create a Session object
 
         Parameters
         ----------
 
         debug : bool
             whether to trigger extra debugging statements
         packer/unpacker : str : 'json', 'pickle' or import_string
             importstrings for methods to serialize message parts.  If just
             'json' or 'pickle', predefined JSON and pickle packers will be used.
             Otherwise, the entire importstring must be used.
 
             The functions must accept at least valid JSON input, and output
             *bytes*.
 
             For example, to use msgpack:
             packer = 'msgpack.packb', unpacker='msgpack.unpackb'
         pack/unpack : callables
             You can also set the pack/unpack callables for serialization
             directly.
         session : unicode (must be ascii)
             the ID of this Session object.  The default is to generate a new
             UUID.
         bsession : bytes
             The session as bytes
         username : unicode
             username added to message headers.  The default is to ask the OS.
         key : bytes
             The key used to initialize an HMAC signature.  If unset, messages
             will not be signed or checked.
         signature_scheme : str
             The message digest scheme. Currently must be of the form 'hmac-HASH',
             where 'HASH' is a hashing function available in Python's hashlib.
             The default is 'hmac-sha256'.
             This is ignored if 'key' is empty.
         keyfile : filepath
             The file containing a key.  If this is set, `key` will be
             initialized to the contents of the file.
         """
         super(Session, self).__init__(**kwargs)
         self._check_packers()
         self.none = self.pack({})
         # ensure self._session_default() if necessary, so bsession is defined:
         self.session
         self.pid = os.getpid()
         self._new_auth()
 
     @property
     def msg_id(self):
         """always return new uuid"""
         return str(uuid.uuid4())
 
     def _check_packers(self):
         """check packers for datetime support."""
         pack = self.pack
         unpack = self.unpack
 
         # check simple serialization
         msg = dict(a=[1,'hi'])
         try:
             packed = pack(msg)
         except Exception as e:
             msg = "packer '{packer}' could not serialize a simple message: {e}{jsonmsg}"
             if self.packer == 'json':
                 jsonmsg = "\nzmq.utils.jsonapi.jsonmod = %s" % jsonapi.jsonmod
             else:
                 jsonmsg = ""
             raise ValueError(
                 msg.format(packer=self.packer, e=e, jsonmsg=jsonmsg)
             )
 
         # ensure packed message is bytes
         if not isinstance(packed, bytes):
             raise ValueError("message packed to %r, but bytes are required"%type(packed))
 
         # check that unpack is pack's inverse
         try:
             unpacked = unpack(packed)
             assert unpacked == msg
         except Exception as e:
             msg = "unpacker '{unpacker}' could not handle output from packer '{packer}': {e}{jsonmsg}"
             if self.packer == 'json':
                 jsonmsg = "\nzmq.utils.jsonapi.jsonmod = %s" % jsonapi.jsonmod
             else:
                 jsonmsg = ""
             raise ValueError(
                 msg.format(packer=self.packer, unpacker=self.unpacker, e=e, jsonmsg=jsonmsg)
             )
 
         # check datetime support
         msg = dict(t=datetime.now())
         try:
             unpacked = unpack(pack(msg))
             if isinstance(unpacked['t'], datetime):
                 raise ValueError("Shouldn't deserialize to datetime")
         except Exception:
             self.pack = lambda o: pack(squash_dates(o))
             self.unpack = lambda s: unpack(s)
 
     def msg_header(self, msg_type):
         return msg_header(self.msg_id, msg_type, self.username, self.session)
 
     def msg(self, msg_type, content=None, parent=None, header=None, metadata=None):
         """Return the nested message dict.
 
         This format is different from what is sent over the wire. The
         serialize/deserialize methods converts this nested message dict to the wire
         format, which is a list of message parts.
         """
         msg = {}
         header = self.msg_header(msg_type) if header is None else header
         msg['header'] = header
         msg['msg_id'] = header['msg_id']
         msg['msg_type'] = header['msg_type']
         msg['parent_header'] = {} if parent is None else extract_header(parent)
         msg['content'] = {} if content is None else content
         msg['metadata'] = self.metadata.copy()
         if metadata is not None:
             msg['metadata'].update(metadata)
         return msg
 
     def sign(self, msg_list):
         """Sign a message with HMAC digest. If no auth, return b''.
 
         Parameters
         ----------
         msg_list : list
             The [p_header,p_parent,p_content] part of the message list.
         """
         if self.auth is None:
             return b''
         h = self.auth.copy()
         for m in msg_list:
             h.update(m)
         return str_to_bytes(h.hexdigest())
 
     def serialize(self, msg, ident=None):
         """Serialize the message components to bytes.
 
         This is roughly the inverse of deserialize. The serialize/deserialize
         methods work with full message lists, whereas pack/unpack work with
         the individual message parts in the message list.
 
         Parameters
         ----------
         msg : dict or Message
             The next message dict as returned by the self.msg method.
 
         Returns
         -------
         msg_list : list
             The list of bytes objects to be sent with the format::
 
                 [ident1, ident2, ..., DELIM, HMAC, p_header, p_parent,
                  p_metadata, p_content, buffer1, buffer2, ...]
 
             In this list, the ``p_*`` entities are the packed or serialized
             versions, so if JSON is used, these are utf8 encoded JSON strings.
         """
         content = msg.get('content', {})
         if content is None:
             content = self.none
         elif isinstance(content, dict):
             content = self.pack(content)
         elif isinstance(content, bytes):
             # content is already packed, as in a relayed message
             pass
         elif isinstance(content, unicode_type):
             # should be bytes, but JSON often spits out unicode
             content = content.encode('utf8')
         else:
             raise TypeError("Content incorrect type: %s"%type(content))
 
         real_message = [self.pack(msg['header']),
                         self.pack(msg['parent_header']),
                         self.pack(msg['metadata']),
                         content,
         ]
 
         to_send = []
 
         if isinstance(ident, list):
             # accept list of idents
             to_send.extend(ident)
         elif ident is not None:
             to_send.append(ident)
         to_send.append(DELIM)
 
         signature = self.sign(real_message)
         to_send.append(signature)
 
         to_send.extend(real_message)
 
         return to_send
 
     def send(self, stream, msg_or_type, content=None, parent=None, ident=None,
              buffers=None, track=False, header=None, metadata=None):
         """Build and send a message via stream or socket.
 
         The message format used by this function internally is as follows:
 
         [ident1,ident2,...,DELIM,HMAC,p_header,p_parent,p_content,
          buffer1,buffer2,...]
 
         The serialize/deserialize methods convert the nested message dict into this
         format.
 
         Parameters
         ----------
 
         stream : zmq.Socket or ZMQStream
             The socket-like object used to send the data.
         msg_or_type : str or Message/dict
             Normally, msg_or_type will be a msg_type unless a message is being
             sent more than once. If a header is supplied, this can be set to
             None and the msg_type will be pulled from the header.
 
         content : dict or None
             The content of the message (ignored if msg_or_type is a message).
         header : dict or None
             The header dict for the message (ignored if msg_to_type is a message).
         parent : Message or dict or None
             The parent or parent header describing the parent of this message
             (ignored if msg_or_type is a message).
         ident : bytes or list of bytes
             The zmq.IDENTITY routing path.
         metadata : dict or None
             The metadata describing the message
         buffers : list or None
             The already-serialized buffers to be appended to the message.
         track : bool
             Whether to track.  Only for use with Sockets, because ZMQStream
             objects cannot track messages.
 
 
         Returns
         -------
         msg : dict
             The constructed message.
         """
         if not isinstance(stream, zmq.Socket):
             # ZMQStreams and dummy sockets do not support tracking.
             track = False
 
         if isinstance(msg_or_type, (Message, dict)):
             # We got a Message or message dict, not a msg_type so don't
             # build a new Message.
             msg = msg_or_type
             buffers = buffers or msg.get('buffers', [])
         else:
             msg = self.msg(msg_or_type, content=content, parent=parent,
                            header=header, metadata=metadata)
         if not os.getpid() == self.pid:
             io.rprint("WARNING: attempted to send message from fork")
             io.rprint(msg)
             return
         buffers = [] if buffers is None else buffers
         if self.adapt_version:
             msg = adapt(msg, self.adapt_version)
         to_send = self.serialize(msg, ident)
         to_send.extend(buffers)
         longest = max([ len(s) for s in to_send ])
         copy = (longest < self.copy_threshold)
 
         if buffers and track and not copy:
             # only really track when we are doing zero-copy buffers
             tracker = stream.send_multipart(to_send, copy=False, track=True)
         else:
             # use dummy tracker, which will be done immediately
             tracker = DONE
             stream.send_multipart(to_send, copy=copy)
 
         if self.debug:
             pprint.pprint(msg)
             pprint.pprint(to_send)
             pprint.pprint(buffers)
 
         msg['tracker'] = tracker
 
         return msg
 
     def send_raw(self, stream, msg_list, flags=0, copy=True, ident=None):
         """Send a raw message via ident path.
 
         This method is used to send a already serialized message.
 
         Parameters
         ----------
         stream : ZMQStream or Socket
             The ZMQ stream or socket to use for sending the message.
         msg_list : list
             The serialized list of messages to send. This only includes the
             [p_header,p_parent,p_metadata,p_content,buffer1,buffer2,...] portion of
             the message.
         ident : ident or list
             A single ident or a list of idents to use in sending.
         """
         to_send = []
         if isinstance(ident, bytes):
             ident = [ident]
         if ident is not None:
             to_send.extend(ident)
 
         to_send.append(DELIM)
         to_send.append(self.sign(msg_list))
         to_send.extend(msg_list)
         stream.send_multipart(to_send, flags, copy=copy)
 
     def recv(self, socket, mode=zmq.NOBLOCK, content=True, copy=True):
         """Receive and unpack a message.
 
         Parameters
         ----------
         socket : ZMQStream or Socket
             The socket or stream to use in receiving.
 
         Returns
         -------
         [idents], msg
             [idents] is a list of idents and msg is a nested message dict of
             same format as self.msg returns.
         """
         if isinstance(socket, ZMQStream):
             socket = socket.socket
         try:
             msg_list = socket.recv_multipart(mode, copy=copy)
         except zmq.ZMQError as e:
             if e.errno == zmq.EAGAIN:
                 # We can convert EAGAIN to None as we know in this case
                 # recv_multipart won't return None.
                 return None,None
             else:
                 raise
         # split multipart message into identity list and message dict
         # invalid large messages can cause very expensive string comparisons
         idents, msg_list = self.feed_identities(msg_list, copy)
         try:
             return idents, self.deserialize(msg_list, content=content, copy=copy)
         except Exception as e:
             # TODO: handle it
             raise e
 
     def feed_identities(self, msg_list, copy=True):
         """Split the identities from the rest of the message.
 
         Feed until DELIM is reached, then return the prefix as idents and
         remainder as msg_list. This is easily broken by setting an IDENT to DELIM,
         but that would be silly.
 
         Parameters
         ----------
         msg_list : a list of Message or bytes objects
             The message to be split.
         copy : bool
             flag determining whether the arguments are bytes or Messages
 
         Returns
         -------
         (idents, msg_list) : two lists
             idents will always be a list of bytes, each of which is a ZMQ
             identity. msg_list will be a list of bytes or zmq.Messages of the
             form [HMAC,p_header,p_parent,p_content,buffer1,buffer2,...] and
             should be unpackable/unserializable via self.deserialize at this
             point.
         """
         if copy:
             idx = msg_list.index(DELIM)
             return msg_list[:idx], msg_list[idx+1:]
         else:
             failed = True
             for idx,m in enumerate(msg_list):
                 if m.bytes == DELIM:
                     failed = False
                     break
             if failed:
                 raise ValueError("DELIM not in msg_list")
             idents, msg_list = msg_list[:idx], msg_list[idx+1:]
             return [m.bytes for m in idents], msg_list
 
     def _add_digest(self, signature):
         """add a digest to history to protect against replay attacks"""
         if self.digest_history_size == 0:
             # no history, never add digests
             return
 
         self.digest_history.add(signature)
         if len(self.digest_history) > self.digest_history_size:
             # threshold reached, cull 10%
             self._cull_digest_history()
 
     def _cull_digest_history(self):
         """cull the digest history
 
         Removes a randomly selected 10% of the digest history
         """
         current = len(self.digest_history)
         n_to_cull = max(int(current // 10), current - self.digest_history_size)
         if n_to_cull >= current:
             self.digest_history = set()
             return
         to_cull = random.sample(self.digest_history, n_to_cull)
         self.digest_history.difference_update(to_cull)
 
     def deserialize(self, msg_list, content=True, copy=True):
         """Unserialize a msg_list to a nested message dict.
 
         This is roughly the inverse of serialize. The serialize/deserialize
         methods work with full message lists, whereas pack/unpack work with
         the individual message parts in the message list.
 
         Parameters
         ----------
         msg_list : list of bytes or Message objects
             The list of message parts of the form [HMAC,p_header,p_parent,
             p_metadata,p_content,buffer1,buffer2,...].
         content : bool (True)
             Whether to unpack the content dict (True), or leave it packed
             (False).
         copy : bool (True)
             Whether msg_list contains bytes (True) or the non-copying Message
             objects in each place (False).
 
         Returns
         -------
         msg : dict
             The nested message dict with top-level keys [header, parent_header,
             content, buffers].  The buffers are returned as memoryviews.
         """
         minlen = 5
         message = {}
         if not copy:
             # pyzmq didn't copy the first parts of the message, so we'll do it
             for i in range(minlen):
                 msg_list[i] = msg_list[i].bytes
         if self.auth is not None:
             signature = msg_list[0]
             if not signature:
                 raise ValueError("Unsigned Message")
             if signature in self.digest_history:
                 raise ValueError("Duplicate Signature: %r" % signature)
             self._add_digest(signature)
             check = self.sign(msg_list[1:5])
             if not compare_digest(signature, check):
                 raise ValueError("Invalid Signature: %r" % signature)
         if not len(msg_list) >= minlen:
             raise TypeError("malformed message, must have at least %i elements"%minlen)
         header = self.unpack(msg_list[1])
         message['header'] = extract_dates(header)
         message['msg_id'] = header['msg_id']
         message['msg_type'] = header['msg_type']
         message['parent_header'] = extract_dates(self.unpack(msg_list[2]))
         message['metadata'] = self.unpack(msg_list[3])
         if content:
             message['content'] = self.unpack(msg_list[4])
         else:
             message['content'] = msg_list[4]
         buffers = [memoryview(b) for b in msg_list[5:]]
         if buffers and buffers[0].shape is None:
             # force copy to workaround pyzmq #646
             buffers = [memoryview(b.bytes) for b in msg_list[5:]]
         message['buffers'] = buffers
         # adapt to the current version
         return adapt(message)
 
     def unserialize(self, *args, **kwargs):
         warnings.warn(
             "Session.unserialize is deprecated. Use Session.deserialize.",
             DeprecationWarning,
         )
         return self.deserialize(*args, **kwargs)
 
 
 def test_msg2obj():
     am = dict(x=1)
     ao = Message(am)
     assert ao.x == am['x']
 
     am['y'] = dict(z=1)
     ao = Message(am)
     assert ao.y.z == am['y']['z']
 
     k1, k2 = 'y', 'z'
     assert ao[k1][k2] == am[k1][k2]
 
     am2 = dict(ao)
     assert am['x'] == am2['x']
     assert am['y']['z'] == am2['y']['z']