python 捕获其他模块生成的 sys.stdout 和 sys.exit

Question

我有如下两个脚本但是我不能修改foo.py 我在 bar.py 中编写了代码来捕获 foo.py 产生的输出，这样我就可以在输出到终端之前修改输出。

但是，捕获不起作用。我做错了什么？

foo.py
----
import sys


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


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

这是我可以修改的模块：

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)

编辑：

最终解决方案

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

问题是 foo.py 中对 sys.exit() 的调用。

当你在bar.py中调用foo()，然后到达foo()函数的最后一行时，解释器立即退出，bar.py之后的剩余部分未执行对 foo() 的调用。因此，您正确地捕获了标准输出——它只是在程序退出之前没有打印出来。

如果您无法更改的文件确实包含对 sys.exit(1) 的直接调用，并且您想在不退出的情况下处理它，则可以将对 foo() 的调用包装在 try/except 块。但是，根据实际调用 sys.exit(1) 的上下文，使用 sys.excepthook

之类的方法处理它可能更有意义

编辑：

回复：您的评论如下，

如果输出可能是 sys.stdout 或 sys.stderr，您应该考虑其中任何一个。例如，您可以这样做：

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

是——sys.exit() 加注 SystemExit，所以你可以这样做：
```
try:
    # your code here
except SystemExit:
    pass
```
虽然这对我来说有点粗略。通常非零退出代码表示程序某处出错，因此 sys.exit(1) 被故意引发可能不是您想要在绝大多数情况下静默抑制的事情。

Answer 2

这是您可以尝试的方法：

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

出于演示目的，我将 foo.py 放在同一个文件中，但主要部分是重定向标准输出

python 捕获其他模块生成的 sys.stdout 和 sys.exit

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

python

capture

sys

output

编辑：