+
+
+class RecordStdout(contextlib.AbstractContextManager):
+ """
+ Record what is emitted to stdout.
+
+ >>> with RecordStdout() as record:
+ ... print("This is a test!")
+ >>> print({record().readline()})
+ {'This is a test!\\n'}
+ >>> record().close()
+ """
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.destination = tempfile.SpooledTemporaryFile(mode='r+')
+ self.recorder: Optional[contextlib.redirect_stdout] = None
+
+ def __enter__(self) -> Callable[[], tempfile.SpooledTemporaryFile]:
+ self.recorder = contextlib.redirect_stdout(self.destination)
+ assert self.recorder is not None
+ self.recorder.__enter__()
+ return lambda: self.destination
+
+ def __exit__(self, *args) -> Literal[False]:
+ assert self.recorder is not None
+ self.recorder.__exit__(*args)
+ self.destination.seek(0)
+ return False
+
+
+class RecordStderr(contextlib.AbstractContextManager):
+ """
+ Record what is emitted to stderr.
+
+ >>> import sys
+ >>> with RecordStderr() as record:
+ ... print("This is a test!", file=sys.stderr)
+ >>> print({record().readline()})
+ {'This is a test!\\n'}
+ >>> record().close()
+ """
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.destination = tempfile.SpooledTemporaryFile(mode='r+')
+ self.recorder: Optional[contextlib.redirect_stdout[Any]] = None
+
+ def __enter__(self) -> Callable[[], tempfile.SpooledTemporaryFile]:
+ self.recorder = contextlib.redirect_stderr(self.destination) # type: ignore
+ assert self.recorder is not None
+ self.recorder.__enter__()
+ return lambda: self.destination
+
+ def __exit__(self, *args) -> Literal[False]:
+ assert self.recorder is not None
+ self.recorder.__exit__(*args)
+ self.destination.seek(0)
+ return False
+
+
+class RecordMultipleStreams(contextlib.AbstractContextManager):
+ """
+ Record the output to more than one stream.
+ """
+
+ def __init__(self, *files) -> None:
+ super().__init__()
+ self.files = [*files]
+ self.destination = tempfile.SpooledTemporaryFile(mode='r+')
+ self.saved_writes: List[Callable[..., Any]] = []
+
+ def __enter__(self) -> Callable[[], tempfile.SpooledTemporaryFile]:
+ for f in self.files:
+ self.saved_writes.append(f.write)
+ f.write = self.destination.write
+ return lambda: self.destination
+
+ def __exit__(self, *args) -> Literal[False]:
+ for f in self.files:
+ f.write = self.saved_writes.pop()
+ self.destination.seek(0)
+ return False
+
+
+if __name__ == '__main__':
+ import doctest
+
+ doctest.testmod()