blog.grace.moe

commit 73ab5afc70ea989dadfbd715dee0e3bd3dc34afa
parent f70a70ec717164d5b23ccdf82c785e5d7ad1ba58
Author: gracefu <81774659+gracefuu@users.noreply.github.com>
Date:   Wed, 23 Apr 2025 01:44:03 +0800

Update make_utils

Diffstat:
M
make_utils.py
 | 
96
++++++++++++++++++++++++++++++++++++++++++++++++-------------------------------

1 file changed, 59 insertions(+), 37 deletions(-)
diff --git a/make_utils.py b/make_utils.py
@@ -1,3 +1,5 @@
+from asyncio.streams import StreamReader, StreamWriter
+from asyncio.subprocess import Process
 from concurrent.futures import Executor
 from typing import Any, Awaitable, Callable, ParamSpec, TypeVar
 import asyncio
@@ -73,65 +75,85 @@ EchoStderr = 2
 EchoAll = 3
 
 
-async def _exec_reader(istream, ostream, echo: Any = False):
+async def _exec_writer(
+    proc: Process,
+    ostream: StreamWriter,
+    input: bytes | bytearray | memoryview | None = None,
+):
+    if input is not None:
+        ostream.write(input)
+        await ostream.drain()
+    return await proc.wait()
+
+
+async def _exec_reader(istream: StreamReader, ostream=None):
     contents = b""
-    async for chunk in istream:
+    while not istream.at_eof():
+        chunk = await istream.read(4096 * 16)
         contents += chunk
-        if echo:
+        if ostream:
             ostream.write(chunk)
             ostream.flush()
     return contents
 
 
+async def communicate_echo_wait(
+    proc: Process,
+    input: bytes | bytearray | memoryview | None = None,
+    echo: int = EchoNothing,
+) -> ShellResult:
+    stdout, stderr, returncode = await asyncio.gather(
+        _exec_reader(
+            proc.stdout,  # type: ignore
+            sys.stdout.buffer if echo & EchoStdout else None,
+        ),
+        _exec_reader(
+            proc.stderr,  # type: ignore
+            sys.stderr.buffer if echo & EchoStderr else None,
+        ),
+        _exec_writer(
+            proc,
+            proc.stdin,  # type: ignore
+            input,
+        ),
+    )
+    return ShellResult(stdout, stderr, returncode)
+
+
 async def exec(
     program,
     *args,
     input: bytes | bytearray | memoryview | None = None,
     echo: int = EchoNothing,
 ) -> ShellResult:
-
-    proc = await asyncio.create_subprocess_exec(
-        program,
-        *args,
-        stdin=subprocess.PIPE,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
+    return await communicate_echo_wait(
+        await asyncio.create_subprocess_exec(
+            program,
+            *args,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        ),
+        input,
+        echo,
     )
 
-    if input is not None:
-        proc.stdin.write(input)  # type: ignore
-        _, stdout, stderr, returncode = await asyncio.gather(
-            proc.stdin.drain(),  # type: ignore
-            _exec_reader(proc.stdout, sys.stdout.buffer, echo=echo & EchoStdout),
-            _exec_reader(proc.stderr, sys.stderr.buffer, echo=echo & EchoStderr),
-            proc.wait(),
-        )
-    else:
-        stdout, stderr, returncode = await asyncio.gather(
-            _exec_reader(proc.stdout, sys.stdout.buffer, echo=echo & EchoStdout),
-            _exec_reader(proc.stderr, sys.stderr.buffer, echo=echo & EchoStderr),
-            proc.wait(),
-        )
-
-    return ShellResult(stdout, stderr, returncode)
-
 
 async def shell(
     cmd,
     input: bytes | bytearray | memoryview | None = None,
     echo: int = EchoNothing,
 ) -> ShellResult:
-    proc = await asyncio.create_subprocess_shell(
-        cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    return await communicate_echo_wait(
+        await asyncio.create_subprocess_shell(
+            cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        ),
+        input,
+        echo,
     )
-    stdout, stderr = await proc.communicate(input)
-    if echo & EchoStdout:
-        sys.stdout.buffer.write(stdout)
-        sys.stdout.buffer.flush()
-    if echo & EchoStderr:
-        sys.stderr.buffer.write(stderr)
-        sys.stderr.buffer.flush()
-    return ShellResult(stdout, stderr, proc.returncode)
 
 
 async def make_main(globals, default_target="all()"):