cli: Fix crash when subprocess has a very long stdout/stderr line

This commit is contained in:
Nuno Campos
2024-07-21 12:16:37 -07:00
parent 610b6cc78c
commit b81612c292
2 changed files with 27 additions and 6 deletions
+26 -5
View File
@@ -131,21 +131,42 @@ async def monitor_stream(
if collect:
ba = bytearray()
def handle(line: bytes):
def handle(line: bytes, overrun: bool):
nonlocal on_line
nonlocal display
if display:
sys.stdout.buffer.write(line)
if overrun:
return
if collect:
ba.extend(line)
if display:
sys.stdout.write(line.decode())
if on_line:
if on_line(line.decode()):
on_line = None
display = True
async for line in stream:
await asyncio.to_thread(handle, line)
"""Adpated from asyncio.StreamReader.readline() to handle LimitOverrunError."""
sep = b"\n"
seplen = len(sep)
while True:
try:
line = await stream.readuntil(sep)
overrun = False
except asyncio.IncompleteReadError as e:
line = e.partial
overrun = False
except asyncio.LimitOverrunError as e:
if stream._buffer.startswith(sep, e.consumed):
line = stream._buffer[: e.consumed + seplen]
else:
line = stream._buffer.clear()
overrun = True
stream._maybe_resume_transport()
await asyncio.to_thread(handle, line, overrun)
if line == b"":
break
if collect:
return ba
else: