1
# -*- coding: utf-8 -*-
3
"""The hook for new user interfaces to take control of progress bar and status
4
display is pass setupRootUiContext an instance of their UI context class, with
5
the same methods as the _Context class defined here.
7
Long-running functions can be decorated with @display_wrap, and will then be
8
given the extra argument 'ui'. 'ui' is a ProgressContext instance with methods
9
.series(), .imap(), .map() and .display(), any one of which will cause a
10
progress-bar to be displayed.
13
def long_running_function(..., ui)
14
ui.display(msg, progress) # progress is between 0.0 and 1.0
16
for item in ui.map(items, function)
18
for item in ui.imap(items, function)
20
for item in ui.series(items)
23
from __future__ import with_statement, division
24
import sys, time, contextlib, functools, warnings
27
from cogent.util import parallel, terminal
29
__author__ = "Peter Maxwell"
30
__copyright__ = "Copyright 2009, The Cogent Project"
31
__credits__ = ["Peter Maxwell"]
36
curses_terminal = terminal.CursesOutput()
37
except terminal.TerminalUnavailableError:
38
curses_terminal = None
40
CODES = curses_terminal.getCodes()
41
bar_template = CODES['GREEN'] + '%s' + CODES['NORMAL'] + '%s'
43
CLEAR = CODES['UP'] + BOL + CODES['CLEAR_EOL']
46
class TextBuffer(object):
47
"""A file-like object which accumulates written text. Specialised for
48
output to a curses terminal in that it uses CLEAR and re-writing to extend
49
incomplete lines instead of just outputting or buffering them. That
50
allows the output to always end at a newline, ready for a progress bar
51
to be shown, without postponing output of any incomplete last line."""
55
self.pending_eol = False
57
def write(self, text):
58
self.chunks.append(text)
60
def regurgitate(self, out):
62
text = ''.join(self.chunks)
65
#out.write(CODES['YELLOW'])
67
if text.endswith('\n'):
68
self.pending_eol = False
71
self.pending_eol = True
72
self.chunks = [text.split('\n')[-1]]
74
#out.write(CODES['NORMAL'])
77
class ProgressContext(object):
78
"""The interface by which cogent algorithms can report progress to the
79
user interface. Calls self.progress_bar.set(progress, message)"""
81
def __init__(self, progress_bar=None, prefix=None, base=0.0, segment=1.0,
82
parent=None, rate=1.0):
83
self.progress_bar = progress_bar
86
self.segment = segment
94
assert progress_bar is parent.progress_bar
95
self.depth = parent.depth + 1
97
self.t_last = parent.t_last
99
self.prefix = prefix or []
100
self.message = self.prefix + [self.msg]
101
self._max_text_len = 0
105
def subcontext(self):
106
"""For any sub-task which may want to report its own progress, but should
107
not get its own progress bar."""
108
if self.depth == self.max_depth:
110
return ProgressContext(
111
progress_bar = self.progress_bar,
112
prefix = self.message,
113
base = self.base+self.progress*self.segment,
114
segment = self.current*self.segment,
118
def display(self, msg=None, progress=None, current=0.0):
119
"""Inform the UI that we are are at 'progress' of the way through and
120
will be doing 'msg' until we reach and report at progress+current.
126
if progress is not None:
127
self.progress = min(progress, 1.0)
130
if current is not None:
131
self.current = current
134
if msg is not None and msg != self.msg:
135
self.msg = self.message[-1] = msg
139
(self.depth==0 and self.progress in [0.0, 1.0]) or
140
time.time() > self.t_last + self.rate):
144
self.progress_bar.set(self.base+self.progress*self.segment, self.message[0])
145
self.t_last = time.time()
149
self.progress_bar.done()
151
# Not much point while cogent is still full of print statements, but
152
# .info() (and maybe other logging analogues such as .warning()) would
153
# avoid the need to capture stdout:
155
#def info(self, text):
156
# """Display some information which may be more than fleetingly useful,
157
# such as a summary of intermediate statistics or a very mild warning.
158
# A GUI should make this information retrievable but not intrusive.
159
# For terminal UIs this is equivalent to printing"""
160
# raise NotImplementedError
162
def series(self, items, noun='', labels=None, start=None, end=1.0):
163
"""Wrap a looped-over list with a progress bar"""
164
if not hasattr(items, '__len__'):
168
step = (end-start) / len(items)
170
assert len(labels) == len(items)
171
elif len(items) == 1:
177
template = '%s%%%sd/%s' % (noun, len(str(goal)), goal)
178
labels = [template % i for i in range(0, len(items))]
179
for (i, item) in enumerate(items):
180
self.display(msg=labels[i], progress=start+step*i, current=step)
182
self.display(progress=end, current=0)
184
def imap(self, f, s, labels=None, **kw):
185
"""Like itertools.imap() but with a progress bar"""
186
with parallel.mpi_split(len(s)) as comm:
187
(size, rank) = (comm.Get_size(), comm.Get_rank())
188
ordinals = range(0, len(s), size)
189
labels = labels and labels[0::size]
190
for start in self.series(ordinals, labels=labels, **kw):
191
chunk = s[start:start+size]
192
if rank < len(chunk):
193
local_result = f(chunk[rank])
196
for result in comm.allgather(local_result)[:len(chunk)]:
199
def eager_map(self, f, s, **kw):
200
"""Like regular Python2 map() but with a progress bar"""
201
return list(self.imap(f,s, **kw))
203
def map(self, f, s, **kw):
204
"""Synonym for eager_map, unlike in Python3"""
205
return self.eager_map(f, s, **kw)
208
class NullContext(ProgressContext):
209
"""A UI context which discards all output. Useful on secondary MPI cpus,
210
and other situations where all output is suppressed"""
211
def subcontext(self, *args, **kw):
214
def display(self, *args, **kw):
221
class LogFileOutput(object):
222
"""A fake progress bar for when progress bars are impossible"""
224
self.t0 = time.time()
226
self.output = sys.stdout # sys.stderr
231
def set(self, progress, message):
233
delta = '+%s' % int(time.time() - self.t0)
234
progress = int(100*progress+0.5)
235
print >>self.output, "%s %5s %3i%% %s" % (
236
self.lpad, delta, progress, message)
239
class CursesTerminalProgressBar(object):
240
"""Wraps stdout and stderr, displaying a progress bar via simple
241
ascii/curses art and scheduling other output around its redraws."""
243
global curses_terminal
244
assert curses_terminal is not None
245
self.curses_terminal = curses_terminal
246
self.stdout = sys.stdout
247
self.stderr = sys.stderr
248
self.stdout_log = TextBuffer()
249
self.stderr_log = TextBuffer()
252
self.pending_eol = False
254
(sys.stdout, sys.stderr, self._stdout, self._stderr) = (
255
self.stdout_log, self.stderr_log, sys.stdout, sys.stderr)
259
(sys.stdout, sys.stderr) = (self._stdout, self._stderr)
261
def set(self, progress, message):
262
"""Clear the existing progress bar, write out any accumulated
263
stdout and stderr, then draw the updated progress bar."""
264
cols = self.curses_terminal.getColumns()
266
if progress is not None:
267
assert 0.0 <= progress <= 1.0, progress
269
dots = int(progress * width)
270
bar = bar_template % (BLOCK * dots, BLOCK * (width-dots))
273
self.stderr.write(CLEAR * (self.line_count))
275
self.stderr.write(BOL)
276
self.stdout_log.regurgitate(self.stdout)
277
self.stderr_log.regurgitate(self.stderr)
279
if progress is not None:
280
self.stderr.writelines([bar, '\n'])
281
if message is not None:
282
self.stderr.writelines([message[:width], '\n'])
283
self.line_count = (progress is not None) + (message is not None)
286
NULL_CONTEXT = NullContext()
287
CURRENT = threading.local()
288
CURRENT.context = None
290
class RootProgressContext(object):
291
"""The context between long running jobs, when there is no progress bar"""
293
def __init__(self, pbar_constructor, rate):
294
self.pbar_constructor = pbar_constructor
297
def subcontext(self):
298
pbar = self.pbar_constructor()
299
return ProgressContext(pbar, rate=self.rate)
302
def setupRootUiContext(progressBarConstructor=None, rate=None):
303
"""Select a UI Context type depending on system environment"""
304
if parallel.getCommunicator().Get_rank() != 0:
306
elif progressBarConstructor is not None:
307
klass = progressBarConstructor
308
elif curses_terminal and sys.stdout.isatty():
309
klass = CursesTerminalProgressBar
310
elif isinstance(sys.stdout, file):
311
klass = LogFileOutput
318
CURRENT.context = NULL_CONTEXT
322
CURRENT.context = RootProgressContext(klass, rate)
325
def display_wrap(slow_function):
326
"""Decorator which give the function its own UI context.
327
The function will receive an extra argument, 'ui',
328
which is used to report progress etc."""
329
@functools.wraps(slow_function)
331
if getattr(CURRENT, 'context', None) is None:
333
parent = CURRENT.context
334
show_progress = kw.pop('show_progress', None)
335
if show_progress is False:
336
# PendingDeprecationWarning?
337
subcontext = NULL_CONTEXT
339
subcontext = parent.subcontext()
340
kw['ui'] = CURRENT.context = subcontext
342
result = slow_function(*args, **kw)
344
CURRENT.context = parent
351
for j in ui.series(range(10)):
357
print "non-linebuffered output, tricky but look:",
358
for i in ui.series(range(10)):
361
print '\nhalfway through, a new line: ',
367
if __name__ == '__main__':
368
#setupRootUiContext(rate=0.2)
371
# This messes up interactive shells a bit:
373
#atexit.register(CURRENT.done)
added
removed
cogent/util/progress_display.py
