Decode output incrementally, with a configurable decoder

linters/cspell/words.txt

@@ -31,3 +31,5 @@ pythonpath
 nonblock
 pathlib
 anyio
+multibyte
+haia

src/runner.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import codecs
 import os
 import platform
 import shlex
@@ -8,6 +9,7 @@
 import sys
 import tempfile
 import time
+from collections.abc import Callable
 from typing import TYPE_CHECKING
 
 from src.errors import ShpyxInternalError, ShpyxOSNotSupportedError, ShpyxVerificationError
@@ -24,6 +26,21 @@
     import fcntl
 
 
+# A callable that returns a fresh incremental decoder for a single command output stream.
+_DecoderFactory = Callable[[], codecs.IncrementalDecoder]
+
+
+def _default_decoder_factory() -> codecs.IncrementalDecoder:
+    """
+    Get the default output decoder factory:
+        1. Decodes incrementally, so a valid multibyte character split across two reads (which are not aligned to
+        character boundaries) is held until the next read completes it, rather than raising.
+        2. Uses UTF-8.
+        3. Replaces genuinely invalid bytes (e.g. binary data) with the Unicode replacement character instead of raising.
+    """
+    return codecs.getincrementaldecoder("utf-8")(errors="replace")
+
+
 def _is_action_required(*, user: bool | None, default: bool) -> bool:
     """
     Returns whether an action needs to be done, based on whether the user required it and the default value of the
@@ -53,6 +70,7 @@ def __init__(
         verify_return_code: bool = True,
         verify_stderr: bool = False,
         use_signal_names: bool = True,
+        decoder_factory: _DecoderFactory = _default_decoder_factory,
     ) -> None:
         """
         Create a command runner.
@@ -67,12 +85,14 @@ def __init__(
             verify_stderr: Whether to raise an exception if anything was written to stderr during the execution.
             use_signal_names:  Whether to log the name of the signal corresponding to a non-zero error code,
                                in case of result verification failure.
+            decoder_factory: Callable that returns a fresh incremental decoder, used to decode the command output.
         """
         self._log_cmd = log_cmd
         self._log_output = log_output
         self._verify_return_code = verify_return_code
         self._verify_stderr = verify_stderr
         self._use_signal_names = use_signal_names
+        self._decoder_factory = decoder_factory
 
     @staticmethod
     def _log(msg: str) -> None:
@@ -88,6 +108,8 @@ def _add_stdout(
         result: ShellCmdResult,
         data: bytes | None,
         log_output: bool | None,
+        decoder: codecs.IncrementalDecoder,
+        final: bool,
     ) -> None:
         """
         Add partial stdout output to the result.
@@ -96,11 +118,13 @@ def _add_stdout(
             result: The result object of the command.
             data: The partial stdout output to add.
             log_output: Whether to log the output, as supplied to `.run`.
+            decoder: Stream decoder.
+            final: Whether this is the last chunk, flushing any trailing incomplete bytes.
         """
-        if not data:
+        decoded_data = decoder.decode(data or b"", final)
+        if not decoded_data:
             return
 
-        decoded_data = data.decode()
         result.stdout += decoded_data
         result.all_output += decoded_data
 
@@ -113,6 +137,8 @@ def _add_stderr(
         result: ShellCmdResult,
         data: bytes | None,
         log_output: bool | None,
+        decoder: codecs.IncrementalDecoder,
+        final: bool,
     ) -> None:
         """
         Add partial stderr output to the result.
@@ -121,11 +147,13 @@ def _add_stderr(
             result: The result object of the command.
             data: The partial stderr output to add.
             log_output: Whether to log the output, as supplied to `.run`.
+            decoder: Stream decoder.
+            final: Whether this is the last chunk, flushing any trailing incomplete bytes.
         """
-        if not data:
+        decoded_data = decoder.decode(data or b"", final)
+        if not decoded_data:
             return
 
-        decoded_data = data.decode()
         result.stderr += decoded_data
         result.all_output += decoded_data
 
@@ -195,6 +223,7 @@ def run(
         env: dict[str, str] | None = None,
         exec_dir: Path | str | None = None,
         unix_raw: bool = False,
+        decoder_factory: _DecoderFactory | None = None,
     ) -> ShellCmdResult:
         """
         Run a shell command.
@@ -228,6 +257,8 @@ def run(
                       This allows capturing all characters from the command output, including cursor movement and
                       colors. This can be useful when the command is an interactive shell, like `psql`.
                       Runner default: `False`.
+            decoder_factory: Callable that returns a fresh incremental decoder, used to decode the command output.
+                             Runner default: `None`, which uses the default decoder factory.
 
         Returns:
             The result, as a `ShellCmdResult` object.
@@ -295,6 +326,11 @@ def run(
         # Initialize the result object.
         result = ShellCmdResult(cmd=cmd_str)
 
+        # Create a fresh decoder per stream (they are stateful and must not be shared).
+        decoder_factory = decoder_factory or self._decoder_factory
+        stdout_decoder = decoder_factory()
+        stderr_decoder = decoder_factory()
+
         # Make all the command outputs non-blocking, so that it can be interrupted.
         if _SYSTEM != "Windows":
             fcntl.fcntl(
@@ -315,15 +351,19 @@ def run(
             stderr_data = p.stderr.read()
 
             # Add partial outputs to result and log them, if needed.
-            self._add_stdout(result=result, data=stdout_data, log_output=log_output)
-            self._add_stderr(result=result, data=stderr_data, log_output=log_output)
+            self._add_stdout(
+                result=result, decoder=stdout_decoder, data=stdout_data, log_output=log_output, final=False
+            )
+            self._add_stderr(
+                result=result, decoder=stderr_decoder, data=stderr_data, log_output=log_output, final=False
+            )
 
             time.sleep(0.01)
 
-        # Get the remaining outputs and add them to the result.
+        # Get the remaining outputs and add them to the result, flushing any trailing incomplete bytes.
         final_stdout, final_stderr = p.communicate()
-        self._add_stdout(result=result, data=final_stdout, log_output=log_output)
-        self._add_stderr(result=result, data=final_stderr, log_output=log_output)
+        self._add_stdout(result=result, decoder=stdout_decoder, data=final_stdout, log_output=log_output, final=True)
+        self._add_stderr(result=result, decoder=stderr_decoder, data=final_stderr, log_output=log_output, final=True)
 
         # Cleanup.
         p.stdout.close()

tests/fake_proc.py

@@ -0,0 +1,43 @@
+import pytest_mock
+
+
+class _FakeStream:
+    """A stdout/stderr stand-in that hands out queued byte chunks one read at a time."""
+
+    def __init__(self, chunks: list[bytes]) -> None:
+        self.chunks = list(chunks)
+
+    def read(self) -> bytes:
+        return self.chunks.pop(0) if self.chunks else b""
+
+    def fileno(self) -> int:
+        return 0
+
+    def close(self) -> None:
+        pass
+
+
+class _FakeProc:
+    """A `subprocess.Popen` stand-in that emits controlled output chunks through the read loop."""
+
+    def __init__(self, stdout_chunks: list[bytes], stderr_chunks: list[bytes]) -> None:
+        self.stdout = _FakeStream(stdout_chunks)
+        self.stderr = _FakeStream(stderr_chunks)
+        self.returncode = 0
+
+    def poll(self) -> int | None:
+        # Keep the read loop going while either stream still has queued chunks.
+        return None if (self.stdout.chunks or self.stderr.chunks) else 0
+
+    def communicate(self) -> tuple[bytes, bytes]:
+        return b"", b""
+
+
+def patch_fake_proc(
+    mocker: pytest_mock.MockerFixture,
+    *,
+    stdout_chunks: list[bytes],
+    stderr_chunks: list[bytes],
+) -> None:
+    proc = _FakeProc(stdout_chunks, stderr_chunks or [])
+    mocker.patch("src.runner.subprocess.Popen", return_value=proc)

tests/test_runner.py

@@ -2,13 +2,21 @@
 Test the default runner, `shpyx.run`.
 """
 
+from __future__ import annotations
+
+import codecs
 import platform
 import signal
 import subprocess
 import tempfile
 from enum import Enum, auto
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+from tests.fake_proc import patch_fake_proc
+
+if TYPE_CHECKING:
+    from _typeshed import ReadableBuffer
 
 import pytest
 import pytest_mock
@@ -305,3 +313,35 @@ def test_unix_raw_enabled() -> None:
     )
     assert result.return_code == 123
     assert result.all_output == stderr_by_platform[_SYSTEM]
+
+
+def test_output_decoding(mocker: pytest_mock.MockerFixture) -> None:
+    """
+    Decoding must gracefully handle two separate hazards in a single run:
+      1. A valid multibyte UTF-8 character split across two output stream reads.
+      2. A genuinely invalid UTF-8 byte in the output (e.g. binary/Latin-1 data).
+    """
+    # '€' is b"\xe2\x82\xac". Split it across two reads, then feed a lone invalid byte (b"\xff").
+    patch_fake_proc(mocker, stdout_chunks=[b"\xe2\x82", b"\xac", b"\xff"], stderr_chunks=[])
+
+    result = shpyx.run("dummy_cmd")
+    assert result.stdout == "€�"
+    assert result.stderr == ""
+
+
+def test_output_decoding_custom_decoder(mocker: pytest_mock.MockerFixture) -> None:
+    """
+    Test the `decoder_factory` argument.
+    """
+
+    class _AppendADecoder(codecs.IncrementalDecoder):
+        # Custom decoder adding 'a' to each byte.
+
+        def decode(self, input: ReadableBuffer, final: bool = False) -> str:  # noqa: A002, FBT001, FBT002, ARG002
+            return "".join(f"{byte:c}a" for byte in bytes(input))
+
+    # 'hi' -> 'h','a','i','a'
+    patch_fake_proc(mocker, stdout_chunks=[b"hi"], stderr_chunks=[])
+
+    result = shpyx.run("dummy_cmd", decoder_factory=_AppendADecoder)
+    assert result.stdout == "haia"