remove use of utils.flatten

use `list(chain.from_iterable))` which is equivalent

ipython_kernel/serialize.py

@@ -1,179 +1,181 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 itertools import chain
+
 from IPython.utils.py3compat import PY3, buffer_to_bytes_py2
-from IPython.utils.data import flatten
 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)
+    arg_bufs = list(chain.from_iterable(
+        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)
+    kwarg_bufs = list(chain.from_iterable(
+        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_parallel/client/map.py

@@ -1,129 +1,124 b''
 # encoding: utf-8
 
 """Classes used in scattering and gathering sequences.
 
 Scattering consists of partitioning a sequence and sending the various
 pieces to individual nodes in a cluster.
 """
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 from __future__ import division
 
 import sys
-from itertools import islice
-
-from IPython.utils.data import flatten as utils_flatten
-
+from itertools import islice, chain
 
 numpy = None
 
 def is_array(obj):
     """Is an object a numpy array?
     
     Avoids importing numpy until it is requested
     """
     global numpy
     if 'numpy' not in sys.modules:
         return False
     
     if numpy is None:
         import numpy
     return isinstance(obj, numpy.ndarray)
 
 class Map(object):
     """A class for partitioning a sequence using a map."""
     
     def getPartition(self, seq, p, q, n=None):
         """Returns the pth partition of q partitions of seq.
         
         The length can be specified as `n`,
         otherwise it is the value of `len(seq)`
         """
         n = len(seq) if n is None else n
         # Test for error conditions here
         if p<0 or p>=q:
           raise ValueError("must have 0 <= p <= q, but have p=%s,q=%s" % (p, q))
         
         remainder = n % q
         basesize = n // q
         
         if p < remainder:
             low = p * (basesize + 1)
             high = low + basesize + 1
         else:
             low = p * basesize + remainder
             high = low + basesize
         
         try:
             result = seq[low:high]
         except TypeError:
             # some objects (iterators) can't be sliced,
             # use islice:
             result = list(islice(seq, low, high))
             
         return result
            
     def joinPartitions(self, listOfPartitions):
         return self.concatenate(listOfPartitions)
     
     def concatenate(self, listOfPartitions):
         testObject = listOfPartitions[0]
         # First see if we have a known array type
         if is_array(testObject):
             return numpy.concatenate(listOfPartitions)
         # Next try for Python sequence types
         if isinstance(testObject, (list, tuple)):
-            return utils_flatten(listOfPartitions)
+            return list(chain.from_iterable(listOfPartitions))
         # If we have scalars, just return listOfPartitions
         return listOfPartitions
 
 class RoundRobinMap(Map):
     """Partitions a sequence in a round robin fashion.
     
     This currently does not work!
     """
 
     def getPartition(self, seq, p, q, n=None):
         n = len(seq) if n is None else n
         return seq[p:n:q]
 
     def joinPartitions(self, listOfPartitions):
         testObject = listOfPartitions[0]
         # First see if we have a known array type
         if is_array(testObject):
             return self.flatten_array(listOfPartitions)
         if isinstance(testObject, (list, tuple)):
             return self.flatten_list(listOfPartitions)
         return listOfPartitions
     
     def flatten_array(self, listOfPartitions):
         test = listOfPartitions[0]
         shape = list(test.shape)
         shape[0] = sum([ p.shape[0] for p in listOfPartitions])
         A = numpy.ndarray(shape)
         N = shape[0]
         q = len(listOfPartitions)
         for p,part in enumerate(listOfPartitions):
             A[p:N:q] = part
         return A
     
     def flatten_list(self, listOfPartitions):
         flat = []
         for i in range(len(listOfPartitions[0])):
             flat.extend([ part[i] for part in listOfPartitions if len(part) > i ])
         return flat
 
 def mappable(obj):
     """return whether an object is mappable or not."""
     if isinstance(obj, (tuple,list)):
         return True
     if is_array(obj):
         return True
     return False
 
 dists = {'b':Map,'r':RoundRobinMap}
 
-    
-