Viper_VRTE/subprocess.vp

import t, c
from stdint import *
import platmacro
import memhub
import string

# Windows backend dependencies
import w32.win32base
import w32.win32file
import w32.win32process
import w32.win32sync

# POSIX backend dependencies
import posix


# Subprocess output buffer size
define OUTPUT_BUF_SIZE = 65536

# Module global MBuddy pointer (initialized in main)
_mbuddy: memhub.MemBuddy*


# =============================================================================
# CompletedProcess result class
# =============================================================================

class CompletedProcess:
    """Result class for subprocess.run."""
    args: str           # Original command line
    returncode: int     # Process exit code (0 means success)
    stdout: str         # Standard output (when capture_output=True)
    stderr: str         # Standard error (when capture_output=True)

    def __new__(self, args: str, returncode: int, stdout: str, stderr: str) -> t.CPtr:
        # Heap allocation (avoiding use-after-return)
        return t.CPtr(_mbuddy.alloc(CompletedProcess.__sizeof__()))

    def __init__(self, args: str, returncode: int, stdout: str, stderr: str):
        self.args = args
        self.returncode = returncode
        self.stdout = stdout
        self.stderr = stderr


# =============================================================================
# Entry function
# =============================================================================

def run(args: str, capture_output: bool = True, text: bool = True) -> CompletedProcess*:
    """Run a subprocess synchronously and wait for it to finish.
    The args are the full command line string.
    If capture_output=True, stdout/stderr will be captured. If text=True, the output will be a string (currently it always returns a string).
    Returns a CompletedProcess object."""
    if platmacro.IS_WINDOWS:
        return _run_win32(args, capture_output, text)
    return _run_posix(args, capture_output, text)


# =============================================================================
# Windows backend implementation
# =============================================================================

def _run_win32(args: str, capture_output: bool, text: bool) -> CompletedProcess*:
    """Windows platform subprocess execution (CreateProcessA pipe)"""
    # Copy command line (CreateProcessA requires a modifiable lpCommandLine)
    args_len: sizet = len(args)
    cmd_line: bytes = _mbuddy.alloc(args_len + 1)
    if cmd_line == None:
        return CompletedProcess(args, -1, None, None)
    i: sizet = 0
    for ch in args:
        cmd_line[i] = ch
        i += 1
    cmd_line[args_len] = 0

    # Security attributes (allow handle inheritance)
    sa: w32.win32base.SECURITY_ATTRIBUTES = w32.win32base.SECURITY_ATTRIBUTES()
    sa.nLength = w32.win32base.SECURITY_ATTRIBUTES.__sizeof__()
    sa.lpSecurityDescriptor = None
    sa.bInheritHandle = 1  # TRUE

    # Create stdout pipe
    stdout_read:  w32.win32base.HANDLE = None
    stdout_write: w32.win32base.HANDLE = None
    if capture_output:
        pipe_result: BOOL = w32.win32file.CreatePipe(
            &stdout_read,
            &stdout_write,
            &sa,
            0
        )
        if pipe_result == 0:
            return CompletedProcess(args, -1, None, None)
        # Prevent read end from being inherited by child process
        w32.win32file.SetHandleInformation(stdout_read, w32.win32file.HANDLE_FLAG_INHERIT, 0)

    # Set startup information
    si: w32.win32process.STARTUPINFOA = w32.win32process.STARTUPINFOA()
    si.cb = w32.win32process.STARTUPINFOA.__sizeof__()
    if capture_output:
        si.dwFlags = w32.win32process.STARTF_USESTDHANDLES
        si.hStdOutput = stdout_write
        si.hStdError = stdout_write  # stderr merged to stdout

    pi: w32.win32process.PROCESS_INFORMATION = w32.win32process.PROCESS_INFORMATION()

    # Create process
    # use CREATE_NO_WINDOW when capture_output=True (output to pipe, no console)
    # use 0                when capture_output=False (inherit parent process console, let stdout directly display)
    create_flags: ULONG = w32.win32process.CREATE_NO_WINDOW if capture_output else 0
    create_result: BOOL = w32.win32process.CreateProcessA(
        None,       # lpApplicationName (NULL, parse from command line)
        cmd_line,        # lpCommandLine (modifiable)
        None,       # lpProcessAttributes
        None,       # lpThreadAttributes
        1,               # bInheritHandles = TRUE
        create_flags,    # dwCreationFlags
        None,       # lpEnvironment (inherit from parent process)
        None,       # lpCurrentDirectory (inherit from parent process)
        &si,
        &pi
    )

    if create_result == 0:
        # Create process failed
        if capture_output:
            w32.win32base.CloseHandle(stdout_read)
            w32.win32base.CloseHandle(stdout_write)
        return CompletedProcess(args, -1, None, None)

    # Close parent process's write end (child process already has a copy of it)
    if capture_output:
        w32.win32base.CloseHandle(stdout_write)

    # Read stdout + wait for process to exit
    # Note: must read and wait at the same time, or else deadlock will occur when child process outputs more than pipe buffer size
    # (child process blocks on writing stdout, parent process blocks on WaitForSingleObject)
    stdout_str: str = None
    if capture_output:
        buf_size: ULONG = OUTPUT_BUF_SIZE
        buf: bytes = _mbuddy.alloc(buf_size)
        if buf != None:
            total: ULONG = 0
            while True:
                bytes_read: ULONG = 0
                read_result: BOOL = w32.win32file.ReadFile(
                    stdout_read,
                    buf + total,
                    buf_size - 1 - total,
                    &bytes_read,
                    None
                )
                if read_result == 0 or bytes_read == 0:
                    break
                total += bytes_read
                # Buffer will be full, double it
                if total >= buf_size - 1 - 4096:
                    new_size: ULONG = buf_size * 2
                    new_buf: bytes = _mbuddy.alloc(new_size)
                    if new_buf != None:
                        string.memcpy(new_buf, buf, total)
                        buf = new_buf
                        buf_size = new_size
                    else:
                        # Expand failed, stop reading
                        break
            buf[total] = 0
            stdout_str = str(buf)

    # Wait for process to exit
    w32.win32sync.WaitForSingleObject(pi.hProcess, w32.win32base.INFINITE)

    # Get exit code
    exit_code: ULONG = 0
    w32.win32process.GetExitCodeProcess(pi.hProcess, t.CPtr(&exit_code))

    # Close handles
    w32.win32base.CloseHandle(pi.hProcess)
    w32.win32base.CloseHandle(pi.hThread)
    if capture_output:
        w32.win32base.CloseHandle(stdout_read)

    return CompletedProcess(args, exit_code, stdout_str, None)


# =============================================================================
# POSIX backend implementation
# =============================================================================

def _run_posix(args: str, capture_output: bool, text: bool) -> CompletedProcess*:
    """POSIX platform subprocess execution (fork, execvp, pipe, waitpid)
    Note: Currently this is a stub function.
    POSIX implementation requires pipe/fork/waitpid, which are not available on Windows MinGW.
    It will be completed when testing on a POSIX platform."""
    return CompletedProcess(args, -1, None, None)