Using contextlib to redirect stdout

Let's say you're using a package made by someone else that uses print statements to write progress information, such as scikit-learn. Then you may have started with code that looks like this.

from sklearn.ensemble import RandomForestRegressor
from sklearn.datasets import make_regression
X, y = make_regression(n_features=4, n_informative=2, random_state=0, shuffle=False)

regr = RandomForestRegressor(
    max_depth=2,
    n_estimators=5,
    random_state=0,
    verbose=100 # <---------- this is what makes sklearn print
)

regr.fit(X, y)

If you wanted to use the output of those print statements somehow in Python, you could redirect the default target of print (stdout), to some other place, like an io.StringIO() object.

This is a great use-case for a context manager:

import io
from contextlib import redirect_stdout

from sklearn.ensemble import RandomForestRegressor
from sklearn.datasets import make_regression

X, y = make_regression(n_features=4, n_informative=2, random_state=0, shuffle=False)

regr = RandomForestRegressor(
    max_depth=2,
    n_estimators=5,
    random_state=0,
    verbose=100 # <---------- this is what makes sklearn print
)

with redirect_stdout(io.StringIO()) as f:
    regr.fit(X, y)

You can now access the printed output via:

print(f.get_value())

In general, you want to be mindful of capturing stdout like that because a downstream user might still be interested in seeing it. But it does demonstrate the flexibility of context managers.