RelMon/python/progressbar.py

0001 import os
0002 import sys
0003 import time
0004 import math
0005
0006
0007 #### PROGRESSBAR Classes END ####
0008
0009 try:
0010     from abc import ABCMeta, abstractmethod
0011 except ImportError:
0012     AbstractWidget = object
0013     abstractmethod = lambda fn: fn
0014 else:
0015     AbstractWidget = ABCMeta('AbstractWidget', (object,), {})
0016 class UnknownLength: pass
0017 class Widget(AbstractWidget):
0018     '''The base class for all widgets
0019
0020     The ProgressBar will call the widget's update value when the widget should
0021     be updated. The widget's size may change between calls, but the widget may
0022     display incorrectly if the size changes drastically and repeatedly.
0023
0024     The boolean TIME_SENSITIVE informs the ProgressBar that it should be
0025     updated more often because it is time sensitive.
0026     '''
0027
0028     TIME_SENSITIVE = False
0029     __slots__ = ()
0030
0031     @abstractmethod
0032     def update(self, pbar):
0033         '''Updates the widget.
0034
0035         pbar - a reference to the calling ProgressBar
0036         '''
0037 class Timer(Widget):
0038     'Widget which displays the elapsed seconds.'
0039
0040     __slots__ = ('format',)
0041     TIME_SENSITIVE = True
0042
0043     def __init__(self, format='Elapsed Time: %s'):
0044         self.format = format
0045
0046     @staticmethod
0047     def format_time(seconds):
0048         'Formats time as the string "HH:MM:SS".'
0049
0050         return str(datetime.timedelta(seconds=int(seconds)))
0051
0052
0053     def update(self, pbar):
0054         'Updates the widget to show the elapsed time.'
0055
0056         return self.format % self.format_time(pbar.seconds_elapsed)
0057 class WidgetHFill(Widget):
0058     '''The base class for all variable width widgets.
0059
0060     This widget is much like the \\hfill command in TeX, it will expand to
0061     fill the line. You can use more than one in the same line, and they will
0062     all have the same width, and together will fill the line.
0063     '''
0064
0065     @abstractmethod
0066     def update(self, pbar, width):
0067         '''Updates the widget providing the total width the widget must fill.
0068
0069         pbar - a reference to the calling ProgressBar
0070         width - The total width the widget must fill
0071         '''
0072 class Bar(WidgetHFill):
0073     'A progress bar which stretches to fill the line.'
0074
0075     __slots__ = ('marker', 'left', 'right', 'fill', 'fill_left')
0076
0077     def __init__(self, marker='#', left='|', right='|', fill=' ',
0078                  fill_left=True):
0079         '''Creates a customizable progress bar.
0080
0081         marker - string or updatable object to use as a marker
0082         left - string or updatable object to use as a left border
0083         right - string or updatable object to use as a right border
0084         fill - character to use for the empty part of the progress bar
0085         fill_left - whether to fill from the left or the right
0086         '''
0087         self.marker = marker
0088         self.left = left
0089         self.right = right
0090         self.fill = fill
0091         self.fill_left = fill_left
0092
0093
0094     def update(self, pbar, width):
0095         'Updates the progress bar and its subcomponents'
0096
0097         left, marked, right = (format_updatable(i, pbar) for i in
0098             (self.left, self.marker, self.right))
0099
0100         width -= len(left) + len(right)
0101         # Marked must *always* have length of 1
0102         if pbar.maxval:
0103             marked *= int(pbar.currval / pbar.maxval * width)
0104         else:
0105             marked = ''
0106
0107         if self.fill_left:
0108             return '%s%s%s' % (left, marked.ljust(width, self.fill), right)
0109         else:
0110             return '%s%s%s' % (left, marked.rjust(width, self.fill), right)
0111 class BouncingBar(Bar):
0112     def update(self, pbar, width):
0113         'Updates the progress bar and its subcomponents'
0114
0115         left, marker, right = (format_updatable(i, pbar) for i in
0116             (self.left, self.marker, self.right))
0117
0118         width -= len(left) + len(right)
0119
0120         if pbar.finished: return '%s%s%s' % (left, width * marker, right)
0121
0122         position = int(pbar.currval % (width * 2 - 1))
0123         if position > width: position = width * 2 - position
0124         lpad = self.fill * (position - 1)
0125         rpad = self.fill * (width - len(marker) - len(lpad))
0126
0127         # Swap if we want to bounce the other way
0128         if not self.fill_left: rpad, lpad = lpad, rpad
0129
0130         return '%s%s%s%s%s' % (left, lpad, marker, rpad, right)
0131
0132 class FormatLabel(Timer):
0133     'Displays a formatted label'
0134
0135     mapping = {
0136         'elapsed': ('seconds_elapsed', Timer.format_time),
0137         'finished': ('finished', None),
0138         'last_update': ('last_update_time', None),
0139         'max': ('maxval', None),
0140         'seconds': ('seconds_elapsed', None),
0141         'start': ('start_time', None),
0142         'value': ('currval', None)
0143     }
0144
0145     __slots__ = ('format',)
0146     def __init__(self, format):
0147         self.format = format
0148
0149     def update(self, pbar):
0150         context = {}
0151         for name, (key, transform) in self.mapping.items():
0152             try:
0153                 value = getattr(pbar, key)
0154
0155                 if transform is None:
0156                     context[name] = value
0157                 else:
0158                     context[name] = transform(value)
0159             except: pass
0160
0161         return self.format % context
0162
0163 class ProgressBar(object):
0164     '''The ProgressBar class which updates and prints the bar.
0165
0166     A common way of using it is like:
0167     >>> pbar = ProgressBar().start()
0168     >>> for i in range(100):
0169     ...    # do something
0170     ...    pbar.update(i+1)
0171     ...
0172     >>> pbar.finish()
0173
0174     You can also use a ProgressBar as an iterator:
0175     >>> progress = ProgressBar()
0176     >>> for i in progress(some_iterable):
0177     ...    # do something
0178     ...
0179
0180     Since the progress bar is incredibly customizable you can specify
0181     different widgets of any type in any order. You can even write your own
0182     widgets! However, since there are already a good number of widgets you
0183     should probably play around with them before moving on to create your own
0184     widgets.
0185
0186     The term_width parameter represents the current terminal width. If the
0187     parameter is set to an integer then the progress bar will use that,
0188     otherwise it will attempt to determine the terminal width falling back to
0189     80 columns if the width cannot be determined.
0190
0191     When implementing a widget's update method you are passed a reference to
0192     the current progress bar. As a result, you have access to the
0193     ProgressBar's methods and attributes. Although there is nothing preventing
0194     you from changing the ProgressBar you should treat it as read only.
0195
0196     Useful methods and attributes include (Public API):
0197      - currval: current progress (0 <= currval <= maxval)
0198      - maxval: maximum (and final) value
0199      - finished: True if the bar has finished (reached 100%)
0200      - start_time: the time when start() method of ProgressBar was called
0201      - seconds_elapsed: seconds elapsed since start_time and last call to
0202                         update
0203      - percentage(): progress in percent [0..100]
0204     '''
0205
0206     __slots__ = ('currval', 'fd', 'finished', 'last_update_time',
0207                  'left_justify', 'maxval', 'next_update', 'num_intervals',
0208                  'poll', 'seconds_elapsed', 'signal_set', 'start_time',
0209                  'term_width', 'update_interval', 'widgets', '_time_sensitive',
0210                  '__iterable')
0211
0212     _DEFAULT_MAXVAL = 100
0213     _DEFAULT_TERMSIZE = 80
0214
0215     def __init__(self, maxval=None, widgets=None, term_width=None, poll=1,
0216                  left_justify=True, fd=sys.stderr):
0217         '''Initializes a progress bar with sane defaults'''
0218
0219         self.maxval = maxval
0220         self.widgets = widgets
0221         self.fd = fd
0222         self.left_justify = left_justify
0223
0224         self.signal_set = False
0225         if term_width is not None:
0226             self.term_width = term_width
0227         else:
0228             try:
0229                 self._handle_resize()
0230                 signal.signal(signal.SIGWINCH, self._handle_resize)
0231                 self.signal_set = True
0232             except (SystemExit, KeyboardInterrupt): raise
0233             except:
0234                 self.term_width = self._env_size()
0235
0236         self.__iterable = None
0237         self._update_widgets()
0238         self.currval = 0
0239         self.finished = False
0240         self.last_update_time = None
0241         self.poll = poll
0242         self.seconds_elapsed = 0
0243         self.start_time = None
0244         self.update_interval = 1
0245
0246
0247     def __call__(self, iterable):
0248         'Use a ProgressBar to iterate through an iterable'
0249
0250         try:
0251             self.maxval = len(iterable)
0252         except:
0253             if self.maxval is None:
0254                 self.maxval = UnknownLength
0255
0256         self.__iterable = iter(iterable)
0257         return self
0258
0259
0260     def __iter__(self):
0261         return self
0262
0263
0264     def __next__(self):
0265         try:
0266             value = next(self.__iterable)
0267             if self.start_time is None: self.start()
0268             else: self.update(self.currval + 1)
0269             return value
0270         except StopIteration:
0271             self.finish()
0272             raise
0273
0274
0275     # Create an alias so that Python 2.x won't complain about not being
0276     # an iterator.
0277     next = __next__
0278
0279
0280     def _env_size(self):
0281         'Tries to find the term_width from the environment.'
0282
0283         return int(os.environ.get('COLUMNS', self._DEFAULT_TERMSIZE)) - 1
0284
0285
0286     def _handle_resize(self, signum=None, frame=None):
0287         'Tries to catch resize signals sent from the terminal.'
0288
0289         h, w = array('h', ioctl(self.fd, termios.TIOCGWINSZ, '\0' * 8))[:2]
0290         self.term_width = w
0291
0292
0293     def percentage(self):
0294         'Returns the progress as a percentage.'
0295         return self.currval * 100.0 / self.maxval
0296
0297     percent = property(percentage)
0298
0299
0300     def _format_widgets(self):
0301         result = []
0302         expanding = []
0303         width = self.term_width
0304
0305         for index, widget in enumerate(self.widgets):
0306             if isinstance(widget, WidgetHFill):
0307                 result.append(widget)
0308                 expanding.insert(0, index)
0309             else:
0310                 widget = format_updatable(widget, self)
0311                 result.append(widget)
0312                 width -= len(widget)
0313
0314         count = len(expanding)
0315         while count:
0316             portion = max(int(math.ceil(width * 1. / count)), 0)
0317             index = expanding.pop()
0318             count -= 1
0319
0320             widget = result[index].update(self, portion)
0321             width -= len(widget)
0322             result[index] = widget
0323
0324         return result
0325
0326
0327     def _format_line(self):
0328         'Joins the widgets and justifies the line'
0329
0330         widgets = ''.join(self._format_widgets())
0331
0332         if self.left_justify: return widgets.ljust(self.term_width)
0333         else: return widgets.rjust(self.term_width)
0334
0335
0336     def _need_update(self):
0337         'Returns whether the ProgressBar should redraw the line.'
0338         if self.currval >= self.next_update or self.finished: return True
0339
0340         delta = time.time() - self.last_update_time
0341         return self._time_sensitive and delta > self.poll
0342
0343
0344     def _update_widgets(self):
0345         'Checks all widgets for the time sensitive bit'
0346
0347         self._time_sensitive = any(getattr(w, 'TIME_SENSITIVE', False)
0348             for w in self.widgets)
0349
0350
0351     def update(self, value=None):
0352         'Updates the ProgressBar to a new value.'
0353
0354         if value is not None and value is not UnknownLength:
0355             if (self.maxval is not UnknownLength
0356                 and not 0 <= value <= self.maxval):
0357
0358                 raise ValueError('Value out of range')
0359
0360             self.currval = value
0361
0362
0363         if not self._need_update(): return
0364         if self.start_time is None:
0365             raise RuntimeError('You must call "start" before calling "update"')
0366
0367         now = time.time()
0368         self.seconds_elapsed = now - self.start_time
0369         self.next_update = self.currval + self.update_interval
0370         self.fd.write(self._format_line() + '\r')
0371         self.last_update_time = now
0372
0373
0374     def start(self):
0375         '''Starts measuring time, and prints the bar at 0%.
0376
0377         It returns self so you can use it like this:
0378         >>> pbar = ProgressBar().start()
0379         >>> for i in range(100):
0380         ...    # do something
0381         ...    pbar.update(i+1)
0382         ...
0383         >>> pbar.finish()
0384         '''
0385
0386         if self.maxval is None:
0387             self.maxval = self._DEFAULT_MAXVAL
0388
0389         self.num_intervals = max(100, self.term_width)
0390         self.next_update = 0
0391
0392         if self.maxval is not UnknownLength:
0393             if self.maxval < 0: raise ValueError('Value out of range')
0394             self.update_interval = self.maxval / self.num_intervals
0395
0396
0397         self.start_time = self.last_update_time = time.time()
0398         self.update(0)
0399
0400         return self
0401
0402
0403     def finish(self):
0404         'Puts the ProgressBar bar in the finished state.'
0405
0406         self.finished = True
0407         self.update(self.maxval)
0408         self.fd.write('\n')
0409         if self.signal_set:
0410             signal.signal(signal.SIGWINCH, signal.SIG_DFL)
0411 def format_updatable(updatable, pbar):
0412     if hasattr(updatable, 'update'): return updatable.update(pbar)
0413     else: return updatable
0414 #### PROGRESSBAR Classes END ####
0415
0416 class infinite_iterator(object):
0417     def __init__(self):
0418         self.n = 1
0419     def __iter__(self):
0420         return self
0421     def next(self):
0422         return 1