src/safe_subprocess/__init__.py

import os
import signal
import fcntl
import time
import subprocess
from typing import List

MAX_BYTES_PER_READ = 1024
SLEEP_BETWEEN_READS = 0.1


class Result:
    timeout: int
    exit_code: int
    stdout: str
    stderr: str

    def __init__(self, timeout, exit_code, stdout, stderr):
        self.timeout = timeout
        self.exit_code = exit_code
        self.stdout = stdout
        self.stderr = stderr


def set_nonblocking(reader):
    fd = reader.fileno()
    fl = fcntl.fcntl(fd, fcntl.F_GETFL)
    fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)


def run(
    args: List[str],
    timeout_seconds: int = 15,
    max_output_size: int = 2048,
    env = None,
    cwd: str | None = None
) -> Result:
    """
    Runs the given program with arguments. After the timeout elapses, kills the process
    and all other processes in the process group. Captures at most max_output_size bytes
    of stdout and stderr each, and discards any output beyond that.
    """
    p = subprocess.Popen(
        args,
        env=env,
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        start_new_session=True,
        bufsize=MAX_BYTES_PER_READ,
        cwd=cwd
    )
    set_nonblocking(p.stdout)
    set_nonblocking(p.stderr)

    process_group_id = os.getpgid(p.pid)

    # We sleep for 0.1 seconds in each iteration.
    max_iterations = timeout_seconds * 10
    stdout_saved_bytes = []
    stderr_saved_bytes = []
    stdout_bytes_read = 0
    stderr_bytes_read = 0

    for _ in range(max_iterations):
        this_stdout_read = p.stdout.read(MAX_BYTES_PER_READ)
        this_stderr_read = p.stderr.read(MAX_BYTES_PER_READ)
        # this_stdout_read and this_stderr_read may be None if stdout or stderr
        # are closed. Without these checks, test_close_output fails.
        if this_stdout_read is not None and stdout_bytes_read < max_output_size:
            stdout_saved_bytes.append(this_stdout_read)
            stdout_bytes_read += len(this_stdout_read)
        if this_stderr_read is not None and stderr_bytes_read < max_output_size:
            stderr_saved_bytes.append(this_stderr_read)
            stderr_bytes_read += len(this_stderr_read)
        exit_code = p.poll()
        if exit_code is not None:
            break
        time.sleep(SLEEP_BETWEEN_READS)

    try:
        # Kills the process group. Without this line, test_fork_once fails.
        os.killpg(process_group_id, signal.SIGKILL)
    except ProcessLookupError:
        pass

    timeout = exit_code is None
    exit_code = exit_code if exit_code is not None else -1
    stdout = b"".join(stdout_saved_bytes).decode("utf-8", errors="ignore")
    stderr = b"".join(stderr_saved_bytes).decode("utf-8", errors="ignore")
    return Result(timeout=timeout, exit_code=exit_code, stdout=stdout, stderr=stderr)