utils.py 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. import sys
  2. class CaptureStream:
  3. def __init__(self):
  4. self.contents = ""
  5. def write(self, text):
  6. self.contents += text
  7. def close(self):
  8. pass
  9. def flush(self):
  10. pass
  11. class CaptureOutput:
  12. def __init__(self, hide):
  13. self.hide = hide
  14. self._prev_streams = None
  15. self.stdout, self.stderr = None, None
  16. def __enter__(self):
  17. if self.hide:
  18. self._prev_streams = [
  19. sys.stdout,
  20. sys.stderr,
  21. sys.__stdout__,
  22. sys.__stderr__,
  23. ]
  24. self.stdout = CaptureStream()
  25. self.stderr = CaptureStream()
  26. sys.stdout, sys.__stdout__ = self.stdout, self.stdout
  27. sys.stderr, sys.__stderr__ = self.stderr, self.stderr
  28. return self
  29. def get_hidden_stdout(self):
  30. return self.stdout.contents
  31. def get_hidden_stderr(self):
  32. return self.stderr.contents
  33. def __exit__(self, exc_type, exc_val, exc_tb):
  34. if self.hide:
  35. sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__ = self._prev_streams
  36. del self._prev_streams