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)

paxton4416 · Accepted Answer · 2021-05-27T08:36:48.543

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,

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()

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.

I can not modify foo and this is what i am trying to capture https://github.com/ansible/ansible/blob/ae50d0518223bb5f261d024eaf79f42304cdca44/lib/ansible/utils/display.py#L284 — Dariusz Krynicki, May 27 '21 at 08:18
is there a way to capture sys.exit called in other module and actualy not exit but handle it? — Dariusz Krynicki, May 27 '21 at 08:24

score 0 · Answer 2 · answered May 27 '21 at 08:31

This is what you could try:

import sys
import os
import importlib

# foo.py

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


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


# bar.py

original = sys.stdout

def start_capture():
    sys.stdout = open('tmpfile', 'w')

def stop_capture():
    sys.stdout.close()
    sys.stdout = original
    data = ""
    with open('tmpfile', 'r') as f:
        data = f.read()
    os.remove('tmpfile')
    return data


start_capture()
try:
    foo()
except SystemExit:
    print('\nFoo Exit!\n')
print(stop_capture())

For demo purpose i put the foo.py in same file but the main part is redirecting the stdout

thank you for the code as it does the job so I appreciate your contribution but actualy @paxton4416 suits better my needs. — Dariusz Krynicki, May 27 '21 at 08:48

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

2 Answers2

Edit: