python capture sys.stdout and sys.exit produced by other module

Question

I have two scripts as below but I can not modify foo.py I have wrote code in bar.py to capture output produced by foo.py so I can modify the output before outputting it to terminal.

But, capture does not work. What am I doing wrong?

foo.py
----
import sys


def some_func():
    sys.stdout.write("hello")
    sys.stdout.flush()


def foo():
    some_func()
    sys.exit(1)

this is my module I can modify:

bar.py
----

import io
from contextlib import redirect_stdout
from foo import foo



f = io.StringIO()
with redirect_stdout(f):
    foo()
s = f.getvalue()
print(s)

EDIT:

final solution

import io
from contextlib import redirect_stdout
from foo import foo


def capture_output(func):
    f = io.StringIO()
    try:
        with redirect_stdout(f):
            func()
    except SystemExit:
        return f.getvalue()


capture_output(foo)

Answer 1

The problem is the call to sys.exit() in foo.py.

When you call foo() in bar.py, and then get to the last line of the foo() function, the interpreter exits immediately and the rest of bar.py after the call to foo() isn't executed. So you are capturing stdout properly -- it just isn't printed before the program exits.

If the file you can't change actually does contain a direct call to sys.exit(1) and you want to handle it without exiting, you could wrap the call to foo() in a try/except block. However, depending on the context in which sys.exit(1) is actually called, it may make more sense to handle it using something like sys.excepthook

Edit:

re: your comments below,

1. if the output can potentially be going to either sys.stdout or sys.stderr, you should account for either. E.g., you could do:

   from contextlib import redirect_stdout, redirect_stderr
   
   
   tmp_stdout, tmp_stderr = StringIO(), StringIO()
   
   with redirect_stdout(tmp_stdout), redirect_stderr(tmp_stderr):
       foo()
   
   tmp_stdout, tmp_stderr = tmp_stdout.getvalue(), tmp_stderr.getvalue()
   

2. Yes -- sys.exit() raises SystemExit, so you could potentially do:

   try:
       # your code here
   except SystemExit:
       pass
   

   although this feels somewhat sketchy to me. Usually a non-zero exit code indicates something went wrong somewhere in the program, so sys.exit(1) being intentionally raised is probably not something you want to suppress silently in the vast majority of cases.