How to capture print of function of imported script

Question

I have a dataloader function for a neural network which defines a generator object so I can easily pass it to the training and testing functions implemented in tensorflow.

I'm currently writing an evaluation script for my little project and it would really be handy to be able to capture the prints of the dataloader (which i call from the evaluation script) in order to know which dataset I'm currently evaluating - further demo code.

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***

Is there any handy way to get the print of the called script or another way to get the information transferred to the Evaluation.py (without changing the dataloader function output)

Thanks in advance!

EDIT:

For other users - I had to wrap the next call of my generator to achieve the result I wanted - e.g.

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

Answer 1

I suggest trying contextlib.redirect_stdout for this task, usage example

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