output_capture.py
2.51 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
""" class OutputCapture to test what scripts print to stdout """
from __future__ import print_function
import sys
import logging
# python 2/3 version conflict:
if sys.version_info.major <= 2:
from StringIO import StringIO
# reload is a builtin
else:
from io import StringIO
if sys.version_info.minor < 4:
from imp import reload
else:
from importlib import reload
class OutputCapture:
""" context manager that captures stdout
use as follows::
with OutputCapture() as capturer:
run_my_script(some_args)
# either test line-by-line ...
for line in capturer:
some_test(line)
# ...or test all output in one go
some_test(capturer.get_data())
In order to solve issues with old logger instances still remembering closed
StringIO instances as "their" stdout, logging is shutdown and restarted
upon entering this Context Manager. This means that you may have to reload
your module, as well.
"""
def __init__(self):
self.buffer = StringIO()
self.orig_stdout = None
self.data = None
def __enter__(self):
# Avoid problems with old logger instances that still remember an old
# closed StringIO as their sys.stdout
logging.shutdown()
reload(logging)
# replace sys.stdout with own buffer.
self.orig_stdout = sys.stdout
sys.stdout = self.buffer
return self
def __exit__(self, exc_type, exc_value, traceback):
sys.stdout = self.orig_stdout # re-set to original
self.data = self.buffer.getvalue()
self.buffer.close() # close buffer
self.buffer = None
if exc_type: # there has been an error
print('Got error during output capture!')
print('Print captured output and re-raise:')
for line in self.data.splitlines():
print(line.rstrip()) # print output before re-raising
def get_data(self):
""" retrieve all the captured data """
if self.buffer is not None:
return self.buffer.getvalue()
elif self.data is not None:
return self.data
else: # should not be possible
raise RuntimeError('programming error or someone messed with data!')
def __iter__(self):
for line in self.get_data().splitlines():
yield line
def reload_module(self, mod):
""" Wrapper around reload function for different python versions """
return reload(mod)