A context manager for discarding stdout and stderr.

hide_stdio.py

import contextlib
import os
import sys


try:
    from io import StringIO as StdioStream
except ImportError:
    from StringIO import StringIO  # 2.7
    class StdioStream(StringIO):
        def write(self, msg):
            StringIO.write(self, msg.decode())
    del StringIO


_IGNORED = StdioStream()


@contextlib.contextmanager
def hide_stdio(*, fds=False):
    """Swallow stdout and stderr."""
    # See the different approaches below.
    with _hide_stdio() as streams:
        if fds:
            orig_stdout, orig_stderr = streams
            with hide_fd(orig_stdout):
                with hide_fd(orig_stderr):
                    yield streams
        else:
            yield streams


#################################################
# file descriptor helpers

# (see https://github.com/microsoft/vscode-python/pull/13494)
@contextlib.contextmanager
def hide_fd(stream):
    # uh-oh: `stream` (even `sys.stdout`) isn't necessarily
    # a file (see `StringIO`)...
    orig_fd = stream.fileno()
    with os.fdopen(os.dup(orig_fd), "w") as dup_stream:
        dup_fd = dup_stream.fileno()
        #with open(os.devnull, "w") as devnull:
        #    os.dup2(devnull.fileno(), orig_fd)
        #try:
        #    yield orig_fd
        #finally:
        #    os.dup2(dup_fd, orig_fd)
        with open(os.devnull, "w") as devnull:
            # Set the original FD (e.g. 1 for stdout) to /dev/null.
            os.dup2(os.dup(devnull.fileno()), orig_fd)
            try:
                yield orig_fd
            finally:
                # Restore the original FD.
                os.dup2(dup_fd, orig_fd)


@contextlib.contextmanager
def hide_fd(stream):
    with open(os.devnull, "w") as devnull:
        with redirect_fd(stream, devnull) as orig_fd:
            yield orig_fd


@contextlib.contextmanager
def redirect_fd(stream, target):
    stream, name = _resolve_stream(stream)
    orig_fd = _get_fileno(stream)
    if orig_fd is None:
        raise Exception(f'stream {stream!r} does not have a fileno')
    # `os.dup2()` will close `target_fd` when we restore the original
    # so we must use a duplicate FD.
    target_fd = _get_fileno(target)
    if target_fd is None:
        raise Exception(f'target {target!r} does not have a fileno')
    # We make a duplicate so we can restore it later.
    # XXX Do we need to worry about inheritance (e.g. FD_CLOEXEC)?
    # See:
    #  - https://docs.python.org/3/library/os.html#os.dup
    #  - https://docs.python.org/3/library/os.html#fd-inheritance
    #  - https://linux.die.net/man/2/dup
    dup_fd = os.dup(orig_fd)
    try:
        # Set the original FD (e.g. 1 for stdout) to the target.
        os.dup2(target_fd, orig_fd)
        try:
            yield (orig_fd, dup_fd)
        finally:
            # Restore the original FD to the stream.
            os.dup2(dup_fd, orig_fd)
    finally:
        # XXX Ignore failures here?
        os.close(dup_fd)


def _resolve_stream(stream):
    if isinstance(stream, str):
        name = stream
        if stream == 'stdout':
            stream = sys.stdout
        elif stream == 'stderr':
            stream = sys.stderr
        else:
            raise ValueError(f'unsupported stream {stream}')
    else:
        if stream is sys.__stdout__:
            name = 'stdout'
        elif stream is sys.__stderr__:
            name = 'stderr'
        else:
            name = None
    return stream, name


def _get_fileno(file):
    if isinstance(file, int):
        fileno = file
        # XXX Also check against the system's max fileno?
        if fileno < 0:
            raise ValueError(f'invalid fileno, got {fileno}')
        return fileno
    try:
        file_fileno = file.fileno
    except AttributeError:
        # We *could* fail here but instead let the caller decide.
        return None
    return file_fileno()

  
#################################################
# approach #1: using stdlib

@contextlib.contextmanager
def _hide_stdio():
    orig = (sys.stdout, sys.stderr)
    with contextlib.redirect_stdout(_IGNORED):
        with contextlib.redirect_stderr(_IGNORED):
            yield orig


#################################################
# approach #2: using custom redirectors

@contextlib.contextmanager
def _hide_stdio():
    with redirect_stdout(_IGNORED) as orig_stdout:
        with redirect_stderr(_IGNORED) as orig_stderr:
            yield (orig_stdout, orig_stderr)


@contextlib.contextmanager
def redirect_stdout(target, *, fd=False):
    orig = sys.stdout
    sys.stdout = target
    try:
        yield orig
    finally:
        sys.stdout = orig


@contextlib.contextmanager
def redirect_stderr(target, *, fd=False):
    orig = sys.stderr
    sys.stderr = target
    try:
        yield orig
    finally:
        sys.stderr = orig


#################################################
# approach #3: hard-coded

@contextlib.contextmanager
def _hide_stdio():
    orig = (sys.stdout, sys.stderr)
    sys.stdout = _IGNORED
    sys.stderr = _IGNORED
    try:
        yield orig
    finally:
        sys.stdout, sys.stderr = orig


#################################################
# approach #4: using a generic redirector

@contextlib.contextmanager
def _hide_stdio():
    with redirect_stdio('stdout', _IGNORED) as orig_stdout:
        with redirect_stdio('stderr', _IGNORED) as orig_stderr:
            yield (orig_stdout, orig_stderr)


def redirect_stdio(stream, *, fd=False):
    if stream == 'stdout':
        orig = sys.stdout
        sys.stdout = target
    elif stream == 'stderr':
        orig = sys.stderr
        sys.stderr = target
    else:
        raise ValueError(f'unsupported stream {stream}')

    try:
        yield orig
    finally:
        if stream == 'stdout':
            sys.stdout = orig
        elif stream == 'stderr':
            sys.stderr = orig