Search code examples
pythonterminalcommand-lineautomationsubprocess

How to encapsulate running a process in Python with realtime + to-variable output capture and return code capture?


Note: I have read this, this, this, and this question. While there is much useful info, none give an exact answer. My knowledge of the language is limited, so I cannot put together pieces from those answers in a way that fit the needed use case (especially point (4) below).

I am looking for a way to run a process with a given set of arguments in current Python (latest version atm is 3.11), with the following requirements:

  1. The stderr and stdout of the process are displayed in real-time, as they would be if run directly (or through a script) in almost any shell, i.e. bash or PowerShell
  2. Both streams are also separately captured into a string or byte array for accessing later
  3. The return code is captured on process finish
  4. This is encapsualted in a function that simply requires the argument list, and returns an object containing both streams and the return code

All points except for (1) are covered by subprocess.run(args, capture_output=True). So, I need to define a function

def run_and_output_realtime:
    # ???
    return result

Which would allow to change just the first line of code like

run_tool_result = subproccess.run(args, capture_output=True)
if run_tool_result.returncode != 0:
    if "Unauthorized" in run_tool_result.stderr.decode():
        print("Please check authorization")
else:
    if "Duration" in run_tool_result.stdout.decode():
        # parse to get whatever duration was in the output

To run_tool_result = run_and_output_realtime(args), and have all the remaining lines unchanged and working.


Solution

  • The following would hopefully solve your (1). Tested it on Python 3.11.

    import subprocess
    import sys
    from typing import List, Tuple, NamedTuple
    
    class ProcessResult(NamedTuple):
        stdout: bytes
        stderr: bytes
        returncode: int
    
    def run_and_output_realtime(args: List[str]) -> ProcessResult:
        stdout_lines = []
        stderr_lines = []
        
        process = subprocess.Popen(
            args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True,
            bufsize=1,
        )
    
        # Read stdout and stderr line by line
        while process.poll() is None:
            stdout_line = process.stdout.readline()
            stderr_line = process.stderr.readline()
    
            if stdout_line:
                sys.stdout.write(stdout_line)
                sys.stdout.flush()
                stdout_lines.append(stdout_line)
            
            if stderr_line:
                sys.stderr.write(stderr_line)
                sys.stderr.flush()
                stderr_lines.append(stderr_line)
    
        # Capture remaining stdout and stderr lines
        stdout_remaining = process.stdout.readlines()
        stderr_remaining = process.stderr.readlines()
        stdout_lines.extend(stdout_remaining)
        stderr_lines.extend(stderr_remaining)
    
        process.stdout.close()
        process.stderr.close()
        
        stdout = "".join(stdout_lines).encode()
        stderr = "".join(stderr_lines).encode()
        
        return ProcessResult(stdout=stdout, stderr=stderr, returncode=process.returncode)