Use offsets in layout, not lengths

Author: davidbrochart

jupyter_server/base/zmqhandlers.py

@@ -162,9 +162,16 @@ def _on_zmq_reply(self, stream, msg_list):
             self.log.warning("zmq message arrived on closed channel")
             self.close()
             return
+        channel = getattr(stream, "channel", None)
+        offsets = []
+        curr_sum = 0
+        for msg in msg_list:
+            length = len(msg)
+            offsets.append(length + curr_sum)
+            curr_sum += length
         layout = json.dumps({
-            "channel": getattr(stream, "channel", None),
-            "offsets": [0] + [len(msg) for msg in msg_list] + [0],
+            "channel": channel,
+            "offsets": offsets,
         }).encode("utf-8")
         layout_length = len(layout).to_bytes(2, byteorder="little")
         bin_msg = b"".join([layout_length, layout] + msg_list)

jupyter_server/services/kernels/handlers.py

@@ -694,17 +694,13 @@ def on_restart_failed(self):
         self._send_status_message("dead")
 
 
-
 def get_msg_list(msg, offsets):
     i0 = 0
-    i = 1
-    while True:
-        i1 = i0 + offsets[i]
-        if i0 == i1:
-            return
+    for i1 in offsets:
         yield msg[i0:i1]
         i0 = i1
-        i += 1
+    yield msg[i0:]
+
 
 # -----------------------------------------------------------------------------
 # URL to handler mappings