Redirect standard output to variable from method with `sys.stdout.write()`-CodePudding

I want to use an external library and redirect their output, but the library uses a function represented below as msg.

I tested some other stdout methods (all of them are in the minimal example below) and all of them are redirected. I also try to use contextmanager without success.

Also if msg method takes file parameter with sys.stdout, the output will be redirected.

import sys
from contextlib import redirect_stdout
from io import StringIO
from contextlib import contextmanager


def msg(msg, file=sys.stdout):
    file.write(msg)


def msg2(msg):
    sys.stdout.write(msg)


def print_method(msg):
    print(msg)


stream = StringIO()
write_to_stream = redirect_stdout(stream)
with write_to_stream:
    print("print")
    msg("method\n")
    msg2("method 2\n")
    print_method("print method")
    sys.stdout.write("write\n")

print("\nCaught data:")
print(stream.getvalue())

The output of example is:

method

Caught data:
print
method not default
method 2
print method
write

But I want to get a line with method under Caught data: line.

Is there an option to redirect method msg with the file default parameter too?

CodePudding user response：

To force redirection of all output to stdout to a file object, you can override sys.stdout directly with the desired file object:

sys.stdout = stream

To make the override temporary, however, you can use unittest.mock.patch as a context manager:

import sys
from io import StringIO
from unittest.mock import patch

stream = StringIO()

with patch('sys.stdout', stream):
    print('foo')
    sys.stdout.write('bar')

print('Captured output:', stream.getvalue(), sep='\n')

This outputs:

Captured output:
foo
bar

Demo: https://replit.com/@blhsing/JealousKnownScience