00001 import os
00002 import sys
00003 import time
00004 import math
00005
00006
00007
00008
00009 try:
00010 from abc import ABCMeta, abstractmethod
00011 except ImportError:
00012 AbstractWidget = object
00013 abstractmethod = lambda fn: fn
00014 else:
00015 AbstractWidget = ABCMeta('AbstractWidget', (object,), {})
00016 class UnknownLength: pass
00017 class Widget(AbstractWidget):
00018 '''The base class for all widgets
00019
00020 The ProgressBar will call the widget's update value when the widget should
00021 be updated. The widget's size may change between calls, but the widget may
00022 display incorrectly if the size changes drastically and repeatedly.
00023
00024 The boolean TIME_SENSITIVE informs the ProgressBar that it should be
00025 updated more often because it is time sensitive.
00026 '''
00027
00028 TIME_SENSITIVE = False
00029 __slots__ = ()
00030
00031 @abstractmethod
00032 def update(self, pbar):
00033 '''Updates the widget.
00034
00035 pbar - a reference to the calling ProgressBar
00036 '''
00037 class Timer(Widget):
00038 'Widget which displays the elapsed seconds.'
00039
00040 __slots__ = ('format',)
00041 TIME_SENSITIVE = True
00042
00043 def __init__(self, format='Elapsed Time: %s'):
00044 self.format = format
00045
00046 @staticmethod
00047 def format_time(seconds):
00048 'Formats time as the string "HH:MM:SS".'
00049
00050 return str(datetime.timedelta(seconds=int(seconds)))
00051
00052
00053 def update(self, pbar):
00054 'Updates the widget to show the elapsed time.'
00055
00056 return self.format % self.format_time(pbar.seconds_elapsed)
00057 class WidgetHFill(Widget):
00058 '''The base class for all variable width widgets.
00059
00060 This widget is much like the \\hfill command in TeX, it will expand to
00061 fill the line. You can use more than one in the same line, and they will
00062 all have the same width, and together will fill the line.
00063 '''
00064
00065 @abstractmethod
00066 def update(self, pbar, width):
00067 '''Updates the widget providing the total width the widget must fill.
00068
00069 pbar - a reference to the calling ProgressBar
00070 width - The total width the widget must fill
00071 '''
00072 class Bar(WidgetHFill):
00073 'A progress bar which stretches to fill the line.'
00074
00075 __slots__ = ('marker', 'left', 'right', 'fill', 'fill_left')
00076
00077 def __init__(self, marker='#', left='|', right='|', fill=' ',
00078 fill_left=True):
00079 '''Creates a customizable progress bar.
00080
00081 marker - string or updatable object to use as a marker
00082 left - string or updatable object to use as a left border
00083 right - string or updatable object to use as a right border
00084 fill - character to use for the empty part of the progress bar
00085 fill_left - whether to fill from the left or the right
00086 '''
00087 self.marker = marker
00088 self.left = left
00089 self.right = right
00090 self.fill = fill
00091 self.fill_left = fill_left
00092
00093
00094 def update(self, pbar, width):
00095 'Updates the progress bar and its subcomponents'
00096
00097 left, marked, right = (format_updatable(i, pbar) for i in
00098 (self.left, self.marker, self.right))
00099
00100 width -= len(left) + len(right)
00101
00102 if pbar.maxval:
00103 marked *= int(pbar.currval / pbar.maxval * width)
00104 else:
00105 marked = ''
00106
00107 if self.fill_left:
00108 return '%s%s%s' % (left, marked.ljust(width, self.fill), right)
00109 else:
00110 return '%s%s%s' % (left, marked.rjust(width, self.fill), right)
00111 class BouncingBar(Bar):
00112 def update(self, pbar, width):
00113 'Updates the progress bar and its subcomponents'
00114
00115 left, marker, right = (format_updatable(i, pbar) for i in
00116 (self.left, self.marker, self.right))
00117
00118 width -= len(left) + len(right)
00119
00120 if pbar.finished: return '%s%s%s' % (left, width * marker, right)
00121
00122 position = int(pbar.currval % (width * 2 - 1))
00123 if position > width: position = width * 2 - position
00124 lpad = self.fill * (position - 1)
00125 rpad = self.fill * (width - len(marker) - len(lpad))
00126
00127
00128 if not self.fill_left: rpad, lpad = lpad, rpad
00129
00130 return '%s%s%s%s%s' % (left, lpad, marker, rpad, right)
00131
00132 class FormatLabel(Timer):
00133 'Displays a formatted label'
00134
00135 mapping = {
00136 'elapsed': ('seconds_elapsed', Timer.format_time),
00137 'finished': ('finished', None),
00138 'last_update': ('last_update_time', None),
00139 'max': ('maxval', None),
00140 'seconds': ('seconds_elapsed', None),
00141 'start': ('start_time', None),
00142 'value': ('currval', None)
00143 }
00144
00145 __slots__ = ('format',)
00146 def __init__(self, format):
00147 self.format = format
00148
00149 def update(self, pbar):
00150 context = {}
00151 for name, (key, transform) in self.mapping.items():
00152 try:
00153 value = getattr(pbar, key)
00154
00155 if transform is None:
00156 context[name] = value
00157 else:
00158 context[name] = transform(value)
00159 except: pass
00160
00161 return self.format % context
00162
00163 class ProgressBar(object):
00164 '''The ProgressBar class which updates and prints the bar.
00165
00166 A common way of using it is like:
00167 >>> pbar = ProgressBar().start()
00168 >>> for i in range(100):
00169 ... # do something
00170 ... pbar.update(i+1)
00171 ...
00172 >>> pbar.finish()
00173
00174 You can also use a ProgressBar as an iterator:
00175 >>> progress = ProgressBar()
00176 >>> for i in progress(some_iterable):
00177 ... # do something
00178 ...
00179
00180 Since the progress bar is incredibly customizable you can specify
00181 different widgets of any type in any order. You can even write your own
00182 widgets! However, since there are already a good number of widgets you
00183 should probably play around with them before moving on to create your own
00184 widgets.
00185
00186 The term_width parameter represents the current terminal width. If the
00187 parameter is set to an integer then the progress bar will use that,
00188 otherwise it will attempt to determine the terminal width falling back to
00189 80 columns if the width cannot be determined.
00190
00191 When implementing a widget's update method you are passed a reference to
00192 the current progress bar. As a result, you have access to the
00193 ProgressBar's methods and attributes. Although there is nothing preventing
00194 you from changing the ProgressBar you should treat it as read only.
00195
00196 Useful methods and attributes include (Public API):
00197 - currval: current progress (0 <= currval <= maxval)
00198 - maxval: maximum (and final) value
00199 - finished: True if the bar has finished (reached 100%)
00200 - start_time: the time when start() method of ProgressBar was called
00201 - seconds_elapsed: seconds elapsed since start_time and last call to
00202 update
00203 - percentage(): progress in percent [0..100]
00204 '''
00205
00206 __slots__ = ('currval', 'fd', 'finished', 'last_update_time',
00207 'left_justify', 'maxval', 'next_update', 'num_intervals',
00208 'poll', 'seconds_elapsed', 'signal_set', 'start_time',
00209 'term_width', 'update_interval', 'widgets', '_time_sensitive',
00210 '__iterable')
00211
00212 _DEFAULT_MAXVAL = 100
00213 _DEFAULT_TERMSIZE = 80
00214
00215 def __init__(self, maxval=None, widgets=None, term_width=None, poll=1,
00216 left_justify=True, fd=sys.stderr):
00217 '''Initializes a progress bar with sane defaults'''
00218
00219 self.maxval = maxval
00220 self.widgets = widgets
00221 self.fd = fd
00222 self.left_justify = left_justify
00223
00224 self.signal_set = False
00225 if term_width is not None:
00226 self.term_width = term_width
00227 else:
00228 try:
00229 self._handle_resize()
00230 signal.signal(signal.SIGWINCH, self._handle_resize)
00231 self.signal_set = True
00232 except (SystemExit, KeyboardInterrupt): raise
00233 except:
00234 self.term_width = self._env_size()
00235
00236 self.__iterable = None
00237 self._update_widgets()
00238 self.currval = 0
00239 self.finished = False
00240 self.last_update_time = None
00241 self.poll = poll
00242 self.seconds_elapsed = 0
00243 self.start_time = None
00244 self.update_interval = 1
00245
00246
00247 def __call__(self, iterable):
00248 'Use a ProgressBar to iterate through an iterable'
00249
00250 try:
00251 self.maxval = len(iterable)
00252 except:
00253 if self.maxval is None:
00254 self.maxval = UnknownLength
00255
00256 self.__iterable = iter(iterable)
00257 return self
00258
00259
00260 def __iter__(self):
00261 return self
00262
00263
00264 def __next__(self):
00265 try:
00266 value = next(self.__iterable)
00267 if self.start_time is None: self.start()
00268 else: self.update(self.currval + 1)
00269 return value
00270 except StopIteration:
00271 self.finish()
00272 raise
00273
00274
00275
00276
00277 next = __next__
00278
00279
00280 def _env_size(self):
00281 'Tries to find the term_width from the environment.'
00282
00283 return int(os.environ.get('COLUMNS', self._DEFAULT_TERMSIZE)) - 1
00284
00285
00286 def _handle_resize(self, signum=None, frame=None):
00287 'Tries to catch resize signals sent from the terminal.'
00288
00289 h, w = array('h', ioctl(self.fd, termios.TIOCGWINSZ, '\0' * 8))[:2]
00290 self.term_width = w
00291
00292
00293 def percentage(self):
00294 'Returns the progress as a percentage.'
00295 return self.currval * 100.0 / self.maxval
00296
00297 percent = property(percentage)
00298
00299
00300 def _format_widgets(self):
00301 result = []
00302 expanding = []
00303 width = self.term_width
00304
00305 for index, widget in enumerate(self.widgets):
00306 if isinstance(widget, WidgetHFill):
00307 result.append(widget)
00308 expanding.insert(0, index)
00309 else:
00310 widget = format_updatable(widget, self)
00311 result.append(widget)
00312 width -= len(widget)
00313
00314 count = len(expanding)
00315 while count:
00316 portion = max(int(math.ceil(width * 1. / count)), 0)
00317 index = expanding.pop()
00318 count -= 1
00319
00320 widget = result[index].update(self, portion)
00321 width -= len(widget)
00322 result[index] = widget
00323
00324 return result
00325
00326
00327 def _format_line(self):
00328 'Joins the widgets and justifies the line'
00329
00330 widgets = ''.join(self._format_widgets())
00331
00332 if self.left_justify: return widgets.ljust(self.term_width)
00333 else: return widgets.rjust(self.term_width)
00334
00335
00336 def _need_update(self):
00337 'Returns whether the ProgressBar should redraw the line.'
00338 if self.currval >= self.next_update or self.finished: return True
00339
00340 delta = time.time() - self.last_update_time
00341 return self._time_sensitive and delta > self.poll
00342
00343
00344 def _update_widgets(self):
00345 'Checks all widgets for the time sensitive bit'
00346
00347 self._time_sensitive = any(getattr(w, 'TIME_SENSITIVE', False)
00348 for w in self.widgets)
00349
00350
00351 def update(self, value=None):
00352 'Updates the ProgressBar to a new value.'
00353
00354 if value is not None and value is not UnknownLength:
00355 if (self.maxval is not UnknownLength
00356 and not 0 <= value <= self.maxval):
00357
00358 raise ValueError('Value out of range')
00359
00360 self.currval = value
00361
00362
00363 if not self._need_update(): return
00364 if self.start_time is None:
00365 raise RuntimeError('You must call "start" before calling "update"')
00366
00367 now = time.time()
00368 self.seconds_elapsed = now - self.start_time
00369 self.next_update = self.currval + self.update_interval
00370 self.fd.write(self._format_line() + '\r')
00371 self.last_update_time = now
00372
00373
00374 def start(self):
00375 '''Starts measuring time, and prints the bar at 0%.
00376
00377 It returns self so you can use it like this:
00378 >>> pbar = ProgressBar().start()
00379 >>> for i in range(100):
00380 ... # do something
00381 ... pbar.update(i+1)
00382 ...
00383 >>> pbar.finish()
00384 '''
00385
00386 if self.maxval is None:
00387 self.maxval = self._DEFAULT_MAXVAL
00388
00389 self.num_intervals = max(100, self.term_width)
00390 self.next_update = 0
00391
00392 if self.maxval is not UnknownLength:
00393 if self.maxval < 0: raise ValueError('Value out of range')
00394 self.update_interval = self.maxval / self.num_intervals
00395
00396
00397 self.start_time = self.last_update_time = time.time()
00398 self.update(0)
00399
00400 return self
00401
00402
00403 def finish(self):
00404 'Puts the ProgressBar bar in the finished state.'
00405
00406 self.finished = True
00407 self.update(self.maxval)
00408 self.fd.write('\n')
00409 if self.signal_set:
00410 signal.signal(signal.SIGWINCH, signal.SIG_DFL)
00411 def format_updatable(updatable, pbar):
00412 if hasattr(updatable, 'update'): return updatable.update(pbar)
00413 else: return updatable
00414
00415
00416 class infinite_iterator(object):
00417 def __init__(self):
00418 self.n = 1
00419 def __iter__(self):
00420 return self
00421 def next(self):
00422 return 1