Writing a capfd-freidnly context manager that *also* captures stdout file descriptor

[hroncok]

Hey folks,

I found myself in a file descriptor hell.

I have a code that needs to take everything that would otherwise go to stdout and do something else with it. Basically like contextlib.redirect_stdout except I need it to work on file descriptor level, because the thing within might print to stdout in a way that is beyond the Python interpreter.

I came up with this (heavily inspired by pytest's own capdf fixture):

import contextlib, io, os, sys, tempfile

@contextlib.contextmanager
def hook_call():
    """Context manager that records all stdout content (on FD level)
    and prints it to stderr at the end, with a 'HOOK STDOUT: ' prefix."""
    tmpfile = io.TextIOWrapper(
        tempfile.TemporaryFile(buffering=0),
        encoding='utf-8',
        errors='replace',
        write_through=True,
    )

    # stdout FD is 1 in reality, but might be replaced by pytest's capfd
    stdout_fd = sys.stdout.fileno()
    stdout_fd_dup = os.dup(stdout_fd)

    # begin capture
    sys.stdout = tmpfile
    os.dup2(tmpfile.fileno(), stdout_fd)
    try:
        yield
    finally:
        # end capture
        sys.stdout = sys.__stdout__
        os.dup2(stdout_fd_dup, stdout_fd)

        tmpfile.seek(0)  # rewind
        for line in tmpfile.read().splitlines():
            print('HOOK STDOUT:', line, file=sys.stderr)
        tmpfile.close()

When I use it from a script:

def make_a_hook_call():
    with hook_call():
        print('from Python to stdout')
        os.system('echo from echo to stdout')
        print('from Python to stderr', file=sys.stderr)
        os.system('echo from echo to stderr >&2')
    print('from Python to stdout outside of hook_call()')

if __name__ == '__main__':
    make_a_hook_call()

It works as expected:

$ python ttt.py >log
from Python to stderr
from echo to stderr
HOOK STDOUT: from Python to stdout
HOOK STDOUT: from echo to stdout
$ cat log
from Python to stdout outside of hook_call()

However, if I try to test this using the capfd fixture:

def test_make_a_hook_call(capfd):
    make_a_hook_call()
    out, err = capfd.readouterr()
    assert out == 'from Python to stdout outside of hook_call()\n'

It breaks:

    def test_make_a_hook_call(capfd):
        make_a_hook_call()
        out, err = capfd.readouterr()
>       assert out == 'from Python to stdout outside of hook_call()\n'
E       AssertionError: assert 'from echo to...hook_call()\n' == 'from Python ...hook_call()\n'
E         + from echo to stdout
E           from Python to stdout outside of hook_call()

My assumption is that trying to write my own capfd and use that in a test with the pytest's capfd is a bad idea...

But maybe I'm just missing something?

Notes:

• Yes, in the real scenario I need to use capfd in the tests. I could use subprocess instead to run the tested thing as a script but that would overly complicate dependency injection currently used to supply a mocked function.

[hroncok]

OK, thanks for the rubber 🦆

I've put this initially to workaround something failing:

# stdout FD is 1 in reality, but might be replaced by pytest's capfd
stdout_fd = sys.stdout.fileno()

But apparently stdout_fd = 1 makes it work as desired 😆

[hroncok]

I also needed to save the original sys.stdout instead of using sys.__stdout__ to restore it to make the code work correctly when pytested.

Final context manager (imports omitted):

@contextlib.contextmanager
def hook_call():
    """Context manager that records all stdout content (on FD level)
    and prints it to stderr at the end, with a 'HOOK STDOUT: ' prefix."""
    tmpfile = io.TextIOWrapper(
        tempfile.TemporaryFile(buffering=0),
        encoding='utf-8',
        errors='replace',
        write_through=True,
    )

    stdout_fd = 1
    stdout_fd_dup = os.dup(stdout_fd)
    stdout_orig = sys.stdout

    # begin capture
    sys.stdout = tmpfile
    os.dup2(tmpfile.fileno(), stdout_fd)
    try:
        yield
    finally:
        # end capture
        sys.stdout = stdout_orig
        os.dup2(stdout_fd_dup, stdout_fd)

        tmpfile.seek(0)  # rewind
        for line in tmpfile.read().splitlines():
            print('HOOK STDOUT:', line, file=sys.stderr)
        tmpfile.close()

Diff from the question:

@@ -9,9 +9,9 @@
         write_through=True,
     )
 
-    # stdout FD is 1 in reality, but might be replaced by pytest's capfd
-    stdout_fd = sys.stdout.fileno()
+    stdout_fd = 1
     stdout_fd_dup = os.dup(stdout_fd)
+    stdout_orig = sys.stdout
 
     # begin capture
     sys.stdout = tmpfile
@@ -20,7 +20,7 @@
         yield
     finally:
         # end capture
-        sys.stdout = sys.__stdout__
+        sys.stdout = stdout_orig
         os.dup2(stdout_fd_dup, stdout_fd)
 
         tmpfile.seek(0)  # rewind

[RonnyPfannschmidt]

This will explode on execnet /pytest-xdist

[hroncok]

Could you please expand on why? (Luckily I don't need this to work everywhere.)

[RonnyPfannschmidt]

Execnet and xdist do quite a lot of magic with stdout /stderr due to how execnet uses a binary protocol on stdio

[hroncok]

This particular case works for me anyway (with pytest-xdist).