1313epoch = datetime .datetime (year = 1970 , month = 1 , day = 1 )
1414
1515
16+ class WrappingIO :
17+
18+ def __init__ (self , target , capturing = False ):
19+ self .buffer = six .StringIO ()
20+ self .target = target
21+ self .capturing = capturing
22+
23+ def write (self , value ):
24+ if self .capturing :
25+ self .buffer .write (value )
26+ else :
27+ self .target .write (value )
28+
29+ def flush (self ):
30+ self .target .write (self .buffer .getvalue ())
31+ self .buffer .seek (0 )
32+ self .buffer .truncate (0 )
33+
34+
1635class StreamWrapper (object ):
1736 '''Wrap stdout and stderr globally'''
1837
@@ -23,13 +42,32 @@ def __init__(self):
2342 self .wrapped_stdout = 0
2443 self .wrapped_stderr = 0
2544 self .wrapped_excepthook = 0
45+ self .capturing = 0
2646
2747 if os .environ .get ('WRAP_STDOUT' ): # pragma: no cover
2848 self .wrap_stdout ()
2949
3050 if os .environ .get ('WRAP_STDERR' ): # pragma: no cover
3151 self .wrap_stderr ()
3252
53+ def start_capturing (self ):
54+ self .capturing += 1
55+ self .update_capturing ()
56+
57+ def stop_capturing (self ):
58+ self .capturing -= 1
59+ self .update_capturing ()
60+
61+ def update_capturing (self ): # pragma: no cover
62+ if isinstance (self .stdout , WrappingIO ):
63+ self .stdout .capturing = self .capturing > 0
64+
65+ if isinstance (self .stderr , WrappingIO ):
66+ self .stderr .capturing = self .capturing > 0
67+
68+ if self .capturing <= 0 :
69+ self .flush ()
70+
3371 def wrap (self , stdout = False , stderr = False ):
3472 if stdout :
3573 self .wrap_stdout ()
@@ -41,7 +79,7 @@ def wrap_stdout(self):
4179 self .wrap_excepthook ()
4280
4381 if not self .wrapped_stdout :
44- self .stdout = sys .stdout = six . StringIO ( )
82+ self .stdout = sys .stdout = WrappingIO ( self . original_stdout )
4583 self .wrapped_stdout += 1
4684
4785 return sys .stdout
@@ -50,7 +88,7 @@ def wrap_stderr(self):
5088 self .wrap_excepthook ()
5189
5290 if not self .wrapped_stderr :
53- self .stderr = sys .stderr = six . StringIO ( )
91+ self .stderr = sys .stderr = WrappingIO ( self . original_stderr )
5492 self .wrapped_stderr += 1
5593
5694 return sys .stderr
@@ -88,22 +126,18 @@ def unwrap_stderr(self):
88126 self .wrapped_stderr = 0
89127
90128 def flush (self ):
91- if self .wrapped_stdout :
129+ if self .wrapped_stdout : # pragma: no branch
92130 try :
93- self .original_stdout .write (self .stdout .getvalue ())
94- self .stdout .seek (0 )
95- self .stdout .truncate (0 )
131+ self .stdout .flush ()
96132 except (io .UnsupportedOperation ,
97133 AttributeError ): # pragma: no cover
98134 self .wrapped_stdout = False
99135 logger .warn ('Disabling stdout redirection, %r is not seekable' ,
100136 sys .stdout )
101137
102- if self .wrapped_stderr :
138+ if self .wrapped_stderr : # pragma: no branch
103139 try :
104- self .original_stderr .write (self .stderr .getvalue ())
105- self .stderr .seek (0 )
106- self .stderr .truncate (0 )
140+ self .stderr .flush ()
107141 except (io .UnsupportedOperation ,
108142 AttributeError ): # pragma: no cover
109143 self .wrapped_stderr = False
0 commit comments