Skip to content

Commit 4eb44b2

Browse files
committed
making sure output capturing is only enabled when we have a running progress bar
1 parent bbba9f4 commit 4eb44b2

3 files changed

Lines changed: 54 additions & 11 deletions

File tree

README.rst

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,14 @@ environment variable, on Linux/Unix systems this can be done through:
118118
119119
# WRAP_STDERR=true python your_script.py
120120
121+
If you need to flush manually while wrapping, you can do so using:
122+
123+
.. code:: python
124+
125+
import progressbar
126+
127+
progressbar.streams.flush()
128+
121129
In most cases the following will work as well, as long as you initialize the
122130
`StreamHandler` after the wrapping has taken place.
123131

progressbar/bar.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ def start(self, *args, **kwargs):
123123
self.stdout = utils.streams.stdout
124124
self.stderr = utils.streams.stderr
125125

126+
utils.streams.start_capturing()
126127
DefaultFdMixin.start(self, *args, **kwargs)
127128

128129
def update(self, value=None):
@@ -132,7 +133,7 @@ def update(self, value=None):
132133

133134
def finish(self, end='\n'):
134135
DefaultFdMixin.finish(self, end=end)
135-
utils.streams.flush()
136+
utils.streams.stop_capturing()
136137
if self.redirect_stdout:
137138
utils.streams.unwrap_stdout()
138139

progressbar/utils.py

Lines changed: 44 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,25 @@
1313
epoch = 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+
1635
class 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

Comments
 (0)