From a2b2dea964632b4970ca264dcfcc6e1ddee35327 2023-04-18 12:24:35
From: Garland Zhang <garland.zhang@databricks.com>
Date: 2023-04-18 12:24:35
Subject: [PATCH] Read each from stream until separator found or buffer is full

---

diff --git a/IPython/core/magics/script.py b/IPython/core/magics/script.py
index e0615c0..d090f47 100644
--- a/IPython/core/magics/script.py
+++ b/IPython/core/magics/script.py
@@ -4,6 +4,7 @@
 # Distributed under the terms of the Modified BSD License.
 
 import asyncio
+import asyncio.exceptions
 import atexit
 import errno
 import os
@@ -208,15 +209,23 @@ class ScriptMagics(Magics):
             """Call a coroutine on the asyncio thread"""
             return asyncio.run_coroutine_threadsafe(coro, event_loop).result()
 
+        async def _readchunk(stream):
+            try:
+                return await stream.readuntil(b'\n')
+            except asyncio.exceptions.IncompleteReadError as e:
+                return e.partial
+            except asyncio.exceptions.LimitOverrunError as e:
+                return await stream.read(e.consumed)
+
         async def _handle_stream(stream, stream_arg, file_object):
             while True:
-                line = (await stream.readline()).decode("utf8", errors="replace")
-                if not line:
+                chunk = (await _readchunk(stream)).decode("utf8", errors="replace")
+                if not chunk:
                     break
                 if stream_arg:
-                    self.shell.user_ns[stream_arg] = line
+                    self.shell.user_ns[stream_arg] = chunk
                 else:
-                    file_object.write(line)
+                    file_object.write(chunk)
                     file_object.flush()
 
         async def _stream_communicate(process, cell):