Merge pull request #7798 from jasongrout/buffer-memoryview

minrk · minrk · commit 84817ef06084 · 2015-03-20T09:49:27.000-04:00
Extract session buffers as memoryviews
diff --git a/IPython/html/base/zmqhandlers.py b/IPython/html/base/zmqhandlers.py
@@ -8,6 +8,7 @@
 import json
 import struct
 import warnings
+import sys
 
 try:
     from urllib.parse import urlparse # Py 3
@@ -43,6 +44,8 @@ def serialize_binary_message(msg):
     # don't modify msg or buffer list in-place
     msg = msg.copy()
     buffers = list(msg.pop('buffers'))
+    if sys.version_info < (3, 4):
+        buffers = [x.tobytes() for x in buffers]
     bmsg = json.dumps(msg, default=date_default).encode('utf8')
     buffers.insert(0, bmsg)
     nbufs = len(buffers)
diff --git a/IPython/html/tests/test_serialize.py b/IPython/html/tests/test_serialize.py
@@ -13,14 +13,14 @@
 def test_serialize_binary():
     s = Session()
     msg = s.msg('data_pub', content={'a': 'b'})
-    msg['buffers'] = [ os.urandom(3) for i in range(3) ]
+    msg['buffers'] = [ memoryview(os.urandom(3)) for i in range(3) ]
     bmsg = serialize_binary_message(msg)
     nt.assert_is_instance(bmsg, bytes)
 
 def test_deserialize_binary():
     s = Session()
     msg = s.msg('data_pub', content={'a': 'b'})
-    msg['buffers'] = [ os.urandom(2) for i in range(3) ]
+    msg['buffers'] = [ memoryview(os.urandom(2)) for i in range(3) ]
     bmsg = serialize_binary_message(msg)
     msg2 = deserialize_binary_message(bmsg)
     nt.assert_equal(msg2, msg)
diff --git a/IPython/kernel/zmq/session.py b/IPython/kernel/zmq/session.py
@@ -810,18 +810,19 @@ def deserialize(self, msg_list, content=True, copy=True):
             Whether to unpack the content dict (True), or leave it packed
             (False).
         copy : bool (True)
-            Whether to return the bytes (True), or the non-copying Message
-            object in each place (False).
+            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].
+            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:
@@ -846,8 +847,11 @@ def deserialize(self, msg_list, content=True, copy=True):
             message['content'] = self.unpack(msg_list[4])
         else:
             message['content'] = msg_list[4]
-
-        message['buffers'] = msg_list[5:]
+        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)
     
