c-level stdout redirection on windows

redirect.py

""" Tested on Windows 10, 64 bit, Python 3.6

Sources: 
https://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/
https://stackoverflow.com/questions/17942874/stdout-redirection-with-ctypes
""" 
from contextlib import contextmanager
import ctypes
import io
import os, sys
import tempfile


### ALL THIS IS NEW ########################################
if sys.version_info < (3, 5):
    libc = ctypes.CDLL(ctypes.util.find_library('c'))
else:
    if hasattr(sys, 'gettotalrefcount'): # debug build
        libc = ctypes.CDLL('ucrtbased')
    else:
        libc = ctypes.CDLL('api-ms-win-crt-stdio-l1-1-0')
        
# c_stdout = ctypes.c_void_p.in_dll(libc, 'stdout')
kernel32 = ctypes.WinDLL('kernel32')
STD_OUTPUT_HANDLE = -11
c_stdout = kernel32.GetStdHandle(STD_OUTPUT_HANDLE)
##############################################################

@contextmanager
def stdout_redirector(stream):
    # The original fd stdout points to. Usually 1 on POSIX systems.
    original_stdout_fd = sys.stdout.fileno()

    def _redirect_stdout(to_fd):
        """Redirect stdout to the given file descriptor."""
        # Flush the C-level buffer stdout
        libc.fflush(None)   #### CHANGED THIS ARG TO NONE #############
        # Flush and close sys.stdout - also closes the file descriptor (fd)
        sys.stdout.close()
        # Make original_stdout_fd point to the same file as to_fd
        os.dup2(to_fd, original_stdout_fd)
        # Create a new sys.stdout that points to the redirected fd
        sys.stdout = io.TextIOWrapper(os.fdopen(original_stdout_fd, 'wb'))

    # Save a copy of the original stdout fd in saved_stdout_fd
    saved_stdout_fd = os.dup(original_stdout_fd)
    try:
        # Create a temporary file and redirect stdout to it
        tfile = tempfile.TemporaryFile(mode='w+b')
        _redirect_stdout(tfile.fileno())
        # Yield to caller, then redirect stdout back to the saved fd
        yield
        _redirect_stdout(saved_stdout_fd)
        # Copy contents of temporary file to the given stream
        tfile.flush()
        tfile.seek(0, io.SEEK_SET)
        stream.write(tfile.read())
    finally:
        tfile.close()
        os.close(saved_stdout_fd)
        
#### Test it
f = io.BytesIO()

with stdout_redirector(f):
    print('foobar')
    print(12)
    libc.puts(b'this comes from C')
    os.system('echo and this is from echo')
    
print('Got stdout: "{0}"'.format(f.getvalue().decode('utf-8')))

I found this helpful and educational. Thank you!

If anyone wants to use this with Jupyter Notebook, it will most likely not work as Jupyter overwrites stdout with something that doesn't have fileno().
This workaround worked for me:

try:
	original_stdout_fd = sys.stdout.fileno()
except io.UnsupportedOperation: # stdout has been replaced, we fall back onto __stdout__
	original_stdout_fd = sys.__stdout__.fileno()

See __stdout__ doc

Thank you buddy.

One more pitfall with stdout_redirector is in the following code where logger continues to use old sys.stdout:

import sys
from redirect import stdout_redirector

def logger(*args, out=sys.stdout):
    print(*args, file=out)

#### Test it
f = io.BytesIO()
with stdout_redirector(f):
    print('foobar')   # it is OK: print uses updated sys.stdout
    logger('foobar')  # not OK: logger users old sys.stdout object

print('Got stdout: "{0}"'.format(f.getvalue().decode('utf-8')))

logger("HELLO")  # not OK: logger will use old sys.stdout, not the one set by stdout_redirector

And one more pitfall is the following place:

        # Flush and close sys.stdout - also closes the file descriptor (fd)
        sys.stdout.close()
        # Make original_stdout_fd point to the same file as to_fd
        os.dup2(to_fd, original_stdout_fd)

Since sys.stdout.close() also closes the file descriptor then the descriptor becomes immediately available for reuse. Any other code, e.g., another thread or a signal handler, can "occupy" original_stdout_fd before os.dup is called.

To deal with this drawback one can change

        # Create a new sys.stdout that points to the redirected fd
        sys.stdout = io.TextIOWrapper(os.fdopen(original_stdout_fd, 'wb'))

to

        # Create a new sys.stdout that points to the redirected fd
        sys.stdout = io.TextIOWrapper(os.fdopen(original_stdout_fd, 'wb', closefd=False)))

Then sys.stdout.close won't close the file descriptor but instead os.dup2 will close it and reassign atomically (at least dup2 on Linux - I haven't found a respective guarantee mentioned in Windows' documentation).

This is nice but I think it might be safer to just launch your program with output redirection via the shell command launching it. Or write your own code wrapper that does just that (rather than insisting on redirecting from within a running program).