Fix subprocess stream decoding

Author: Jensen Zane

pip_audit/_subprocess.py

@@ -5,7 +5,11 @@
 
 import os.path
 import subprocess
+import threading
+import time
+from codecs import getincrementaldecoder
 from collections.abc import Sequence
+from io import BufferedReader
 from subprocess import Popen
 
 from ._state import AuditState
@@ -24,6 +28,14 @@ def __init__(self, msg: str, *, stderr: str) -> None:
         self.stderr = stderr
 
 
+def _read_stream(stream: BufferedReader, output: bytearray) -> None:
+    """
+    Read a subprocess stream into the given output buffer.
+    """
+    while chunk := stream.read(8192):
+        output.extend(chunk)
+
+
 def run(args: Sequence[str], *, log_stdout: bool = False, state: AuditState = AuditState()) -> str:
     """
     Execute the given arguments.
@@ -39,29 +51,46 @@ def run(args: Sequence[str], *, log_stdout: bool = False, state: AuditState = Au
     # state updates, so we trim the first argument down to its basename.
     pretty_args = " ".join([os.path.basename(args[0]), *args[1:]])
 
-    terminated = False
-    stdout = b""
-    stderr = b""
+    stdout = bytearray()
+    stderr = bytearray()
 
     # Run the process with unbuffered I/O, to make the poll-and-read loop below
     # more responsive.
     with Popen(args, bufsize=0, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as process:
-        # NOTE: We use `poll()` to control this loop instead of the `read()` call
-        # to prevent deadlocks. Similarly, `read(size)` will return an empty bytes
-        # once `stdout` hits EOF, so we don't have to worry about that blocking.
-        while not terminated:
-            terminated = process.poll() is not None
-            stdout += process.stdout.read()  # type: ignore
-            stderr += process.stderr.read()  # type: ignore
+        assert process.stdout is not None
+        assert process.stderr is not None
+
+        stdout_thread = threading.Thread(target=_read_stream, args=(process.stdout, stdout))
+        stderr_thread = threading.Thread(target=_read_stream, args=(process.stderr, stderr))
+        stdout_thread.start()
+        stderr_thread.start()
+
+        stdout_decoder = getincrementaldecoder("utf-8")(errors="replace")
+        stdout_decoded = ""
+        stdout_decoded_len = 0
+
+        while process.poll() is None:
+            stdout_decoded += stdout_decoder.decode(bytes(stdout[stdout_decoded_len:]))
+            stdout_decoded_len = len(stdout)
             state.update_state(
                 f"Running {pretty_args}",
-                stdout.decode(errors="replace") if log_stdout else None,
+                stdout_decoded if log_stdout else None,
             )
+            time.sleep(0.1)
+
+        stdout_thread.join()
+        stderr_thread.join()
+
+        stdout_decoded += stdout_decoder.decode(bytes(stdout[stdout_decoded_len:]), final=True)
+        state.update_state(
+            f"Running {pretty_args}",
+            stdout_decoded if log_stdout else None,
+        )
 
         if process.returncode != 0:
             raise CalledProcessError(
                 f"{pretty_args} exited with {process.returncode}",
-                stderr=stderr.decode(errors="replace"),
+                stderr=stderr.decode("utf-8", errors="replace"),
             )
 
     return stdout.decode("utf-8", errors="replace")

test/test_subprocess.py

@@ -1,8 +1,65 @@
+import sys
+
 import pytest
 
 from pip_audit._subprocess import CalledProcessError, run
 
 
+class RecordingState:
+    def __init__(self):
+        self.logs = []
+
+    def update_state(self, message, logs=None):
+        self.logs.append(logs)
+
+
 def test_run_raises():
     with pytest.raises(CalledProcessError):
         run(["false"])
+
+
+def test_run_handles_split_multibyte_stdout():
+    state = RecordingState()
+
+    stdout = run(
+        [
+            sys.executable,
+            "-c",
+            (
+                "import sys, time; "
+                "data = 'é'.encode(); "
+                "sys.stdout.buffer.write(data[:1]); "
+                "sys.stdout.flush(); "
+                "time.sleep(0.2); "
+                "sys.stdout.buffer.write(data[1:]); "
+                "sys.stdout.flush()"
+            ),
+        ],
+        log_stdout=True,
+        state=state,
+    )
+
+    assert stdout == "é"
+    assert "\ufffd" not in "".join(log or "" for log in state.logs)
+
+
+def test_run_handles_split_multibyte_stderr():
+    with pytest.raises(CalledProcessError) as excinfo:
+        run(
+            [
+                sys.executable,
+                "-c",
+                (
+                    "import sys, time; "
+                    "data = 'é'.encode(); "
+                    "sys.stderr.buffer.write(data[:1]); "
+                    "sys.stderr.flush(); "
+                    "time.sleep(0.2); "
+                    "sys.stderr.buffer.write(data[1:]); "
+                    "sys.stderr.flush(); "
+                    "sys.exit(1)"
+                ),
+            ],
+        )
+
+    assert excinfo.value.stderr == "é"