Source code for schrodinger.test.ioredirect

import io
import os
import sys
import tempfile

import schrodinger.utils.log as log


def replace_output_logger_stream(stream):
    """
    Replace the stream used by output loggers.

    """
    handler = log.get_output_logger_handler()
    handler.flush()
    handler.stream = stream


class IORedirect:
    "Use to redirect stderr & stdout."

    def __init__(self, stdout, stderr):
        self._stdout = stdout
        self._stderr = stderr

    def __enter__(self):
        self.old_stdout, self.old_stderr = sys.stdout, sys.stderr
        sys.stdout.flush()
        sys.stderr.flush()
        sys.stdout, sys.stderr = self._stdout, self._stderr
        replace_output_logger_stream(self._stdout)

    def __exit__(self, exc_type, exc_value, traceback):
        self._stderr.flush(), self._stdout.flush()
        sys.stdout = self.old_stdout
        sys.stderr = self.old_stderr
        replace_output_logger_stream(self.old_stdout)


class IOCapture(IORedirect):
    "Use to redirect stderr & stdout to a StringIO object."

    def __init__(self):
        self.string = io.StringIO()
        super().__init__(self.string, self.string)

    def __enter__(self):
        super().__enter__()
        return self.string


class IOSilence:
    '''
    A context manager for doing a "deep suppression" of stdout and stderr,
    which will point the file descriptors to devnull.

    Useful for cases where C/Fortran code does not use mmerr and there is an
    expected error message.
    '''

    def __init__(self):
        # Open a pair of null files
        self.null_fds = [os.open(os.devnull, os.O_RDWR) for x in range(2)]
        # Save the actual stdout (1) and stderr (2) file descriptors.
        self.save_fds = (os.dup(1), os.dup(2))

    def __enter__(self):
        # Assign the null pointers to stdout and stderr.
        os.dup2(self.null_fds[0], 1)
        os.dup2(self.null_fds[1], 2)

    def __exit__(self, *_):
        # Re-assign the real stdout/stderr back to (1) and (2)
        os.dup2(self.save_fds[0], 1)
        os.dup2(self.save_fds[1], 2)
        # Close the null files
        os.close(self.null_fds[0])
        os.close(self.null_fds[1])
        os.close(self.save_fds[0])
        os.close(self.save_fds[1])


class FullIORedirect(IOSilence):
    """
    Context manager returns an IOString with all output from
    stderr/stdout.

    Because of output from c / python, this output may not be in
    the correct order. This is indeed to be used to verify certain
    messages from errors, in tests only.
    """

    def __init__(self):
        self.tempfile = tempfile.NamedTemporaryFile(delete=False)
        self.save_fds = (os.dup(1), os.dup(2))
        self._stdout = sys.stdout
        self._stderr = sys.stderr
        self.string = io.StringIO()

    def __enter__(self):
        # Assign the null pointers to stdout and stderr.
        sys.stdout.flush()
        sys.stderr.flush()
        os.dup2(self.tempfile.fileno(), 1)
        os.dup2(self.tempfile.fileno(), 2)
        sys.stdout = self.string
        sys.stderr = self.string
        return self.string

    def __exit__(self, *_):
        # Re-assign the real stdout/stderr back to (1) and (2)
        os.dup2(self.save_fds[0], 1)
        os.dup2(self.save_fds[1], 2)
        self.tempfile.close()
        sys.stdout = self._stdout
        sys.stderr = self._stderr
        with open(self.tempfile.name) as fh:
            self.string.writelines(fh.readlines())
        os.remove(self.tempfile.name)
        os.close(self.save_fds[0])
        os.close(self.save_fds[1])