如何捕获导入脚本功能的打印

How to capture print of function of imported script

我有一个用于神经网络的数据加载器函数,它定义了一个生成器对象,因此我可以轻松地将它传递给在 tensorflow 中实现的训练和测试函数。

我目前正在为我的小项目编写一个评估脚本,如果能够捕获数据加载器(我从评估脚本中调用)的打印件以便了解我的数据集,那将非常方便我目前正在评估 - 进一步的演示代码。

Dataloader.py:

def dataloader(batch_size=4, training=true):
    ...
    print(current_dataset)
    ....
    yield input_data, ground_truth

Evaluation.py:

from Dataloader import dataloader

train_gen = dataloader(batsize=5, training=True)

train_data, gt = next(train_gen)
current_dataset = ***print of Dataloader.py***

是否有任何方便的方法来获取被调用脚本的打印或其他方法来将信息传输到 Evaluation.py(不更改数据加载器函数输出)

提前致谢!

编辑:

对于其他用户 - 我必须包装我的生成器的下一次调用以获得我想要的结果 - 例如

from Dataloader import dataloader
import io
import contextlib

train_gen = dataloader(batsize=5, training=True)

with contextlib.redirect_stdout(io.StringIO()) as f:
    train_data, gt = next(train_gen)
current_dataset = f.getvalue()

我建议尝试 contextlib.redirect_stdout 完成此任务,用法示例

import io
import contextlib
def func():
    print("123")
    return 1
with contextlib.redirect_stdout(io.StringIO()) as f:
    x = func()
output = f.getvalue()
print(int(output))  # 123