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() 函数的最后一行时，解释器将立即退出，并且在调用foo() 之后bar.py 的其余部分不会执行.所以你正在正确地捕获标准输出——它只是在程序退出之前没有被打印出来。

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

编辑：

re: 你的 cmets 下面，

1. 如果输出可能是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()
   

2. 是的——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 放在同一个文件中，但主要部分是重定向标准输出