如何编写捕获函数所有标准输出的Python装饰器？

要实现这个需求，核心是临时替换sys.stdout为自定义流对象，让这个对象同时把输出内容转发到原控制台和目标文件（或捕获为字符串）。以下是可直接运行的实现方案：

1. 自定义多输出流类 #

先实现一个类，重写write和flush方法，支持同时向多个输出目标写入内容：

import sys
from io import StringIO

class MultiWriter:
    def __init__(self, *streams):
        self.streams = streams

    def write(self, data):
        # 向所有注册的流写入内容
        for stream in self.streams:
            stream.write(data)
            # 立即刷新，避免输出缓冲延迟
            stream.flush()

    def flush(self):
        for stream in self.streams:
            stream.flush()

2. 实现装饰器 #

装饰器需要完成三个核心操作：保存原sys.stdout、替换为自定义流、执行函数后恢复原流。下面提供两种常用模式的实现：

模式1：同时输出到控制台和日志文件 #

def log_to_file(file_path):
    def decorator(func):
        def wrapper(*args, **kwargs):
            # 保存原控制台输出流
            original_stdout = sys.stdout
            # 以追加模式打开日志文件
            with open(file_path, 'a', encoding='utf-8') as log_file:
                # 替换stdout为MultiWriter，同时输出到控制台和文件
                sys.stdout = MultiWriter(original_stdout, log_file)
                try:
                    # 执行原函数
                    return func(*args, **kwargs)
                finally:
                    # 无论函数是否抛出异常，都恢复原输出流
                    sys.stdout = original_stdout
        return wrapper
    return decorator

模式2：捕获输出为字符串，同时保留控制台打印 #

def capture_output(func):
    def wrapper(*args, **kwargs):
        original_stdout = sys.stdout
        # 创建StringIO对象用于捕获输出内容
        captured_output = StringIO()
        sys.stdout = MultiWriter(original_stdout, captured_output)
        try:
            result = func(*args, **kwargs)
            # 获取捕获到的字符串
            output_str = captured_output.getvalue()
            return result, output_str
        finally:
            sys.stdout = original_stdout
            captured_output.close()
    return wrapper

3. 测试示例 #

# 测试日志文件装饰器
@log_to_file('function_output.log')
def test_print():
    print("这是第一条测试输出")
    print("第二条输出，包含中文")
    for i in range(3):
        print(f"循环输出第{i}次")

test_print()

# 测试捕获字符串装饰器
@capture_output
def another_test():
    print("要被捕获的内容")
    return "函数返回值"

result, output = another_test()
print("\n捕获到的输出内容：")
print(repr(output))

关键注意事项 #

• 异常安全：用try-finally确保无论函数执行是否出错，都会恢复原sys.stdout，避免破坏全局输出流。
• 缓冲处理：在write方法中主动调用flush，避免因为输出缓冲导致内容延迟写入文件或捕获不完整。
• 编码问题：打开文件时指定encoding='utf-8'，避免中文输出出现乱码。

内容的提问来源于stack exchange，提问作者Toothpick Anemone