io_spoof.py

"""IO spoofing."""

import sys
from contextlib import contextmanager
from io import StringIO
from collections import deque


@contextmanager
def capture_stdout():
    """Captures standard out within context block.
    Returns an Out object. stdout so far captured is Out.captured.
    
        with capture_stdout() as out:
            print('foo')
            assert out.captured == 'foo\n'
    """
    oldout = sys.stdout
    newout = StringIO()
    sys.stdout = newout

    class Out:
        @property
        def captured(self):
            try:
                val = newout.getvalue()
                self._captured = val
                return self._captured
            except ValueError:
                # After closing context manager.
                return self._captured
        
        def __str__(self):
            return self.captured

    result = Out()
    try:
        yield result
    finally:
        # Clean up (regardless of an exception).
        result.captured  # set result._captured before closing newout
        newout.close()
        sys.stdout = oldout


class StdinError(Exception): pass

class FakeStdin(object):
    """A queue that can be used as python stdin.
    This is useful for automatically testing a function that takes
    user input with the input() function.
    
        with FakeStdin() as stdin:
            stdin.put('foo')
            assert input() == 'foo'
    """
    def __init__(self, write_input=False):
        self.write_input = write_input
        self._queue = deque()

    def __enter__(self):
        self._oldstdin = sys.stdin
        sys.stdin = self
        return self

    def __exit__(self, *exc):
        sys.stdin = self._oldstdin
        return False

    def put(self, *lines):
        self._queue.extend(lines)

    def readline(self):
        try:
            line = self._queue.popleft()
            if callable(line):
                # This allows something like (lambda: time.sleep(1) or 'foo').
                line = line()
            if self.write_input:
                # We write the line to make it look like someone typed it.
                sys.stdout.write(line + '\n')
            return line
        except IndexError:
            raise StdinError('No stdin available.')

    def clear(self):
        self._queue.clear()


def main():
    with capture_stdout() as out:
        print('foo')
        assert out.captured == 'foo\n'

    with FakeStdin() as stdin:
        stdin.put('foo')
        assert input() == 'foo'


if __name__ == '__main__':
    main()