Patrick Williams | c124f4f | 2015-09-15 14:41:29 -0500 | [diff] [blame] | 1 | #!/usr/bin/env python |
| 2 | # -*- coding: iso-8859-1 -*- |
| 3 | # |
| 4 | # progressbar - Text progressbar library for python. |
| 5 | # Copyright (c) 2005 Nilton Volpato |
| 6 | # |
| 7 | # This library is free software; you can redistribute it and/or |
| 8 | # modify it under the terms of the GNU Lesser General Public |
| 9 | # License as published by the Free Software Foundation; either |
| 10 | # version 2.1 of the License, or (at your option) any later version. |
| 11 | # |
| 12 | # This library is distributed in the hope that it will be useful, |
| 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
| 15 | # Lesser General Public License for more details. |
| 16 | # |
| 17 | # You should have received a copy of the GNU Lesser General Public |
| 18 | # License along with this library; if not, write to the Free Software |
| 19 | # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA |
| 20 | |
| 21 | |
| 22 | """Text progressbar library for python. |
| 23 | |
| 24 | This library provides a text mode progressbar. This is typically used |
| 25 | to display the progress of a long running operation, providing a |
| 26 | visual clue that processing is underway. |
| 27 | |
| 28 | The ProgressBar class manages the progress, and the format of the line |
| 29 | is given by a number of widgets. A widget is an object that may |
| 30 | display diferently depending on the state of the progress. There are |
| 31 | three types of widget: |
| 32 | - a string, which always shows itself; |
| 33 | - a ProgressBarWidget, which may return a diferent value every time |
| 34 | it's update method is called; and |
| 35 | - a ProgressBarWidgetHFill, which is like ProgressBarWidget, except it |
| 36 | expands to fill the remaining width of the line. |
| 37 | |
| 38 | The progressbar module is very easy to use, yet very powerful. And |
| 39 | automatically supports features like auto-resizing when available. |
| 40 | """ |
| 41 | |
| 42 | from __future__ import division |
| 43 | |
| 44 | __author__ = "Nilton Volpato" |
| 45 | __author_email__ = "first-name dot last-name @ gmail.com" |
| 46 | __date__ = "2006-05-07" |
| 47 | __version__ = "2.3-dev" |
| 48 | |
| 49 | import sys, time, os |
| 50 | from array import array |
| 51 | try: |
| 52 | from fcntl import ioctl |
| 53 | import termios |
| 54 | except ImportError: |
| 55 | pass |
| 56 | import signal |
| 57 | try: |
| 58 | basestring |
| 59 | except NameError: |
| 60 | basestring = (str,) |
| 61 | |
| 62 | class ProgressBarWidget(object): |
| 63 | """This is an element of ProgressBar formatting. |
| 64 | |
| 65 | The ProgressBar object will call it's update value when an update |
| 66 | is needed. It's size may change between call, but the results will |
| 67 | not be good if the size changes drastically and repeatedly. |
| 68 | """ |
| 69 | def update(self, pbar): |
| 70 | """Returns the string representing the widget. |
| 71 | |
| 72 | The parameter pbar is a reference to the calling ProgressBar, |
| 73 | where one can access attributes of the class for knowing how |
| 74 | the update must be made. |
| 75 | |
| 76 | At least this function must be overriden.""" |
| 77 | pass |
| 78 | |
| 79 | class ProgressBarWidgetHFill(object): |
| 80 | """This is a variable width element of ProgressBar formatting. |
| 81 | |
| 82 | The ProgressBar object will call it's update value, informing the |
| 83 | width this object must the made. This is like TeX \\hfill, it will |
| 84 | expand to fill the line. You can use more than one in the same |
| 85 | line, and they will all have the same width, and together will |
| 86 | fill the line. |
| 87 | """ |
| 88 | def update(self, pbar, width): |
| 89 | """Returns the string representing the widget. |
| 90 | |
| 91 | The parameter pbar is a reference to the calling ProgressBar, |
| 92 | where one can access attributes of the class for knowing how |
| 93 | the update must be made. The parameter width is the total |
| 94 | horizontal width the widget must have. |
| 95 | |
| 96 | At least this function must be overriden.""" |
| 97 | pass |
| 98 | |
| 99 | |
| 100 | class ETA(ProgressBarWidget): |
| 101 | "Widget for the Estimated Time of Arrival" |
| 102 | def format_time(self, seconds): |
| 103 | return time.strftime('%H:%M:%S', time.gmtime(seconds)) |
| 104 | def update(self, pbar): |
| 105 | if pbar.currval == 0: |
| 106 | return 'ETA: --:--:--' |
| 107 | elif pbar.finished: |
| 108 | return 'Time: %s' % self.format_time(pbar.seconds_elapsed) |
| 109 | else: |
| 110 | elapsed = pbar.seconds_elapsed |
| 111 | eta = elapsed * pbar.maxval / pbar.currval - elapsed |
| 112 | return 'ETA: %s' % self.format_time(eta) |
| 113 | |
| 114 | class FileTransferSpeed(ProgressBarWidget): |
| 115 | "Widget for showing the transfer speed (useful for file transfers)." |
| 116 | def __init__(self, unit='B'): |
| 117 | self.unit = unit |
| 118 | self.fmt = '%6.2f %s' |
| 119 | self.prefixes = ['', 'K', 'M', 'G', 'T', 'P'] |
| 120 | def update(self, pbar): |
| 121 | if pbar.seconds_elapsed < 2e-6:#== 0: |
| 122 | bps = 0.0 |
| 123 | else: |
| 124 | bps = pbar.currval / pbar.seconds_elapsed |
| 125 | spd = bps |
| 126 | for u in self.prefixes: |
| 127 | if spd < 1000: |
| 128 | break |
| 129 | spd /= 1000 |
| 130 | return self.fmt % (spd, u + self.unit + '/s') |
| 131 | |
| 132 | class RotatingMarker(ProgressBarWidget): |
| 133 | "A rotating marker for filling the bar of progress." |
| 134 | def __init__(self, markers='|/-\\'): |
| 135 | self.markers = markers |
| 136 | self.curmark = -1 |
| 137 | def update(self, pbar): |
| 138 | if pbar.finished: |
| 139 | return self.markers[0] |
| 140 | self.curmark = (self.curmark + 1) % len(self.markers) |
| 141 | return self.markers[self.curmark] |
| 142 | |
| 143 | class Percentage(ProgressBarWidget): |
| 144 | "Just the percentage done." |
| 145 | def update(self, pbar): |
| 146 | return '%3d%%' % pbar.percentage() |
| 147 | |
| 148 | class SimpleProgress(ProgressBarWidget): |
| 149 | "Returns what is already done and the total, e.g.: '5 of 47'" |
| 150 | def __init__(self, sep=' of '): |
| 151 | self.sep = sep |
| 152 | def update(self, pbar): |
| 153 | return '%d%s%d' % (pbar.currval, self.sep, pbar.maxval) |
| 154 | |
| 155 | class Bar(ProgressBarWidgetHFill): |
| 156 | "The bar of progress. It will stretch to fill the line." |
| 157 | def __init__(self, marker='#', left='|', right='|'): |
| 158 | self.marker = marker |
| 159 | self.left = left |
| 160 | self.right = right |
| 161 | def _format_marker(self, pbar): |
| 162 | if isinstance(self.marker, basestring): |
| 163 | return self.marker |
| 164 | else: |
| 165 | return self.marker.update(pbar) |
| 166 | def update(self, pbar, width): |
| 167 | percent = pbar.percentage() |
| 168 | cwidth = width - len(self.left) - len(self.right) |
| 169 | marked_width = int(percent * cwidth // 100) |
| 170 | m = self._format_marker(pbar) |
| 171 | bar = (self.left + (m * marked_width).ljust(cwidth) + self.right) |
| 172 | return bar |
| 173 | |
| 174 | class ReverseBar(Bar): |
| 175 | "The reverse bar of progress, or bar of regress. :)" |
| 176 | def update(self, pbar, width): |
| 177 | percent = pbar.percentage() |
| 178 | cwidth = width - len(self.left) - len(self.right) |
| 179 | marked_width = int(percent * cwidth // 100) |
| 180 | m = self._format_marker(pbar) |
| 181 | bar = (self.left + (m*marked_width).rjust(cwidth) + self.right) |
| 182 | return bar |
| 183 | |
| 184 | default_widgets = [Percentage(), ' ', Bar()] |
| 185 | class ProgressBar(object): |
| 186 | """This is the ProgressBar class, it updates and prints the bar. |
| 187 | |
| 188 | A common way of using it is like: |
| 189 | >>> pbar = ProgressBar().start() |
| 190 | >>> for i in xrange(100): |
| 191 | ... # do something |
| 192 | ... pbar.update(i+1) |
| 193 | ... |
| 194 | >>> pbar.finish() |
| 195 | |
| 196 | You can also use a progressbar as an iterator: |
| 197 | >>> progress = ProgressBar() |
| 198 | >>> for i in progress(some_iterable): |
| 199 | ... # do something |
| 200 | ... |
| 201 | |
| 202 | But anything you want to do is possible (well, almost anything). |
| 203 | You can supply different widgets of any type in any order. And you |
| 204 | can even write your own widgets! There are many widgets already |
| 205 | shipped and you should experiment with them. |
| 206 | |
| 207 | The term_width parameter must be an integer or None. In the latter case |
| 208 | it will try to guess it, if it fails it will default to 80 columns. |
| 209 | |
| 210 | When implementing a widget update method you may access any |
| 211 | attribute or function of the ProgressBar object calling the |
| 212 | widget's update method. The most important attributes you would |
| 213 | like to access are: |
| 214 | - currval: current value of the progress, 0 <= currval <= maxval |
| 215 | - maxval: maximum (and final) value of the progress |
| 216 | - finished: True if the bar has finished (reached 100%), False o/w |
| 217 | - start_time: the time when start() method of ProgressBar was called |
| 218 | - seconds_elapsed: seconds elapsed since start_time |
| 219 | - percentage(): percentage of the progress [0..100]. This is a method. |
| 220 | |
| 221 | The attributes above are unlikely to change between different versions, |
| 222 | the other ones may change or cease to exist without notice, so try to rely |
| 223 | only on the ones documented above if you are extending the progress bar. |
| 224 | """ |
| 225 | |
| 226 | __slots__ = ('currval', 'fd', 'finished', 'last_update_time', 'maxval', |
| 227 | 'next_update', 'num_intervals', 'seconds_elapsed', |
| 228 | 'signal_set', 'start_time', 'term_width', 'update_interval', |
| 229 | 'widgets', '_iterable') |
| 230 | |
| 231 | _DEFAULT_MAXVAL = 100 |
| 232 | |
| 233 | def __init__(self, maxval=None, widgets=default_widgets, term_width=None, |
| 234 | fd=sys.stderr): |
| 235 | self.maxval = maxval |
| 236 | self.widgets = widgets |
| 237 | self.fd = fd |
| 238 | self.signal_set = False |
| 239 | if term_width is not None: |
| 240 | self.term_width = term_width |
| 241 | else: |
| 242 | try: |
| 243 | self._handle_resize(None, None) |
| 244 | signal.signal(signal.SIGWINCH, self._handle_resize) |
| 245 | self.signal_set = True |
| 246 | except (SystemExit, KeyboardInterrupt): |
| 247 | raise |
| 248 | except: |
| 249 | self.term_width = int(os.environ.get('COLUMNS', 80)) - 1 |
| 250 | |
| 251 | self.currval = 0 |
| 252 | self.finished = False |
| 253 | self.start_time = None |
| 254 | self.last_update_time = None |
| 255 | self.seconds_elapsed = 0 |
| 256 | self._iterable = None |
| 257 | |
| 258 | def __call__(self, iterable): |
| 259 | try: |
| 260 | self.maxval = len(iterable) |
| 261 | except TypeError: |
| 262 | # If the iterable has no length, then rely on the value provided |
| 263 | # by the user, otherwise fail. |
| 264 | if not (isinstance(self.maxval, (int, long)) and self.maxval > 0): |
| 265 | raise RuntimeError('Could not determine maxval from iterable. ' |
| 266 | 'You must explicitly provide a maxval.') |
| 267 | self._iterable = iter(iterable) |
| 268 | self.start() |
| 269 | return self |
| 270 | |
| 271 | def __iter__(self): |
| 272 | return self |
| 273 | |
| 274 | def next(self): |
| 275 | try: |
| 276 | next = self._iterable.next() |
| 277 | self.update(self.currval + 1) |
| 278 | return next |
| 279 | except StopIteration: |
| 280 | self.finish() |
| 281 | raise |
| 282 | |
| 283 | def _handle_resize(self, signum, frame): |
| 284 | h, w = array('h', ioctl(self.fd, termios.TIOCGWINSZ, '\0' * 8))[:2] |
| 285 | self.term_width = w |
| 286 | |
| 287 | def percentage(self): |
| 288 | "Returns the percentage of the progress." |
| 289 | return self.currval * 100.0 / self.maxval |
| 290 | |
| 291 | def _format_widgets(self): |
| 292 | r = [] |
| 293 | hfill_inds = [] |
| 294 | num_hfill = 0 |
| 295 | currwidth = 0 |
| 296 | for i, w in enumerate(self.widgets): |
| 297 | if isinstance(w, ProgressBarWidgetHFill): |
| 298 | r.append(w) |
| 299 | hfill_inds.append(i) |
| 300 | num_hfill += 1 |
| 301 | elif isinstance(w, basestring): |
| 302 | r.append(w) |
| 303 | currwidth += len(w) |
| 304 | else: |
| 305 | weval = w.update(self) |
| 306 | currwidth += len(weval) |
| 307 | r.append(weval) |
| 308 | for iw in hfill_inds: |
| 309 | widget_width = int((self.term_width - currwidth) // num_hfill) |
| 310 | r[iw] = r[iw].update(self, widget_width) |
| 311 | return r |
| 312 | |
| 313 | def _format_line(self): |
| 314 | return ''.join(self._format_widgets()).ljust(self.term_width) |
| 315 | |
| 316 | def _next_update(self): |
| 317 | return int((int(self.num_intervals * |
| 318 | (self.currval / self.maxval)) + 1) * |
| 319 | self.update_interval) |
| 320 | |
| 321 | def _need_update(self): |
| 322 | """Returns true when the progressbar should print an updated line. |
| 323 | |
| 324 | You can override this method if you want finer grained control over |
| 325 | updates. |
| 326 | |
| 327 | The current implementation is optimized to be as fast as possible and |
| 328 | as economical as possible in the number of updates. However, depending |
| 329 | on your usage you may want to do more updates. For instance, if your |
| 330 | progressbar stays in the same percentage for a long time, and you want |
| 331 | to update other widgets, like ETA, then you could return True after |
| 332 | some time has passed with no updates. |
| 333 | |
| 334 | Ideally you could call self._format_line() and see if it's different |
| 335 | from the previous _format_line() call, but calling _format_line() takes |
| 336 | around 20 times more time than calling this implementation of |
| 337 | _need_update(). |
| 338 | """ |
| 339 | return self.currval >= self.next_update |
| 340 | |
| 341 | def update(self, value): |
| 342 | "Updates the progress bar to a new value." |
| 343 | assert 0 <= value <= self.maxval, '0 <= %d <= %d' % (value, self.maxval) |
| 344 | self.currval = value |
| 345 | if not self._need_update(): |
| 346 | return |
| 347 | if self.start_time is None: |
| 348 | raise RuntimeError('You must call start() before calling update()') |
| 349 | now = time.time() |
| 350 | self.seconds_elapsed = now - self.start_time |
| 351 | self.next_update = self._next_update() |
| 352 | self.fd.write(self._format_line() + '\r') |
| 353 | self.last_update_time = now |
| 354 | |
| 355 | def start(self): |
| 356 | """Starts measuring time, and prints the bar at 0%. |
| 357 | |
| 358 | It returns self so you can use it like this: |
| 359 | >>> pbar = ProgressBar().start() |
| 360 | >>> for i in xrange(100): |
| 361 | ... # do something |
| 362 | ... pbar.update(i+1) |
| 363 | ... |
| 364 | >>> pbar.finish() |
| 365 | """ |
| 366 | if self.maxval is None: |
| 367 | self.maxval = self._DEFAULT_MAXVAL |
| 368 | assert self.maxval > 0 |
| 369 | |
| 370 | self.num_intervals = max(100, self.term_width) |
| 371 | self.update_interval = self.maxval / self.num_intervals |
| 372 | self.next_update = 0 |
| 373 | |
| 374 | self.start_time = self.last_update_time = time.time() |
| 375 | self.update(0) |
| 376 | return self |
| 377 | |
| 378 | def finish(self): |
| 379 | """Used to tell the progress is finished.""" |
| 380 | self.finished = True |
| 381 | self.update(self.maxval) |
| 382 | self.fd.write('\n') |
| 383 | if self.signal_set: |
| 384 | signal.signal(signal.SIGWINCH, signal.SIG_DFL) |