utils/lit/ProgressBar.py

   1 #!/usr/bin/env python
   2
   3 # Source: http://code.activestate.com/recipes/475116/, with
   4 # modifications by Daniel Dunbar.
   5
   6 import sys, re, time
   7
   8 class TerminalController:
   9     """
  10     A class that can be used to portably generate formatted output to
  11     a terminal.
  12
  13     `TerminalController` defines a set of instance variables whose
  14     values are initialized to the control sequence necessary to
  15     perform a given action.  These can be simply included in normal
  16     output to the terminal:
  17
  18         >>> term = TerminalController()
  19         >>> print 'This is '+term.GREEN+'green'+term.NORMAL
  20
  21     Alternatively, the `render()` method can used, which replaces
  22     '${action}' with the string required to perform 'action':
  23
  24         >>> term = TerminalController()
  25         >>> print term.render('This is ${GREEN}green${NORMAL}')
  26
  27     If the terminal doesn't support a given action, then the value of
  28     the corresponding instance variable will be set to ''.  As a
  29     result, the above code will still work on terminals that do not
  30     support color, except that their output will not be colored.
  31     Also, this means that you can test whether the terminal supports a
  32     given action by simply testing the truth value of the
  33     corresponding instance variable:
  34
  35         >>> term = TerminalController()
  36         >>> if term.CLEAR_SCREEN:
  37         ...     print 'This terminal supports clearning the screen.'
  38
  39     Finally, if the width and height of the terminal are known, then
  40     they will be stored in the `COLS` and `LINES` attributes.
  41     """
  42     # Cursor movement:
  43     BOL = ''             #: Move the cursor to the beginning of the line
  44     UP = ''              #: Move the cursor up one line
  45     DOWN = ''            #: Move the cursor down one line
  46     LEFT = ''            #: Move the cursor left one char
  47     RIGHT = ''           #: Move the cursor right one char
  48
  49     # Deletion:
  50     CLEAR_SCREEN = ''    #: Clear the screen and move to home position
  51     CLEAR_EOL = ''       #: Clear to the end of the line.
  52     CLEAR_BOL = ''       #: Clear to the beginning of the line.
  53     CLEAR_EOS = ''       #: Clear to the end of the screen
  54
  55     # Output modes:
  56     BOLD = ''            #: Turn on bold mode
  57     BLINK = ''           #: Turn on blink mode
  58     DIM = ''             #: Turn on half-bright mode
  59     REVERSE = ''         #: Turn on reverse-video mode
  60     NORMAL = ''          #: Turn off all modes
  61
  62     # Cursor display:
  63     HIDE_CURSOR = ''     #: Make the cursor invisible
  64     SHOW_CURSOR = ''     #: Make the cursor visible
  65
  66     # Terminal size:
  67     COLS = None          #: Width of the terminal (None for unknown)
  68     LINES = None         #: Height of the terminal (None for unknown)
  69
  70     # Foreground colors:
  71     BLACK = BLUE = GREEN = CYAN = RED = MAGENTA = YELLOW = WHITE = ''
  72
  73     # Background colors:
  74     BG_BLACK = BG_BLUE = BG_GREEN = BG_CYAN = ''
  75     BG_RED = BG_MAGENTA = BG_YELLOW = BG_WHITE = ''
  76
  77     _STRING_CAPABILITIES = """
  78     BOL=cr UP=cuu1 DOWN=cud1 LEFT=cub1 RIGHT=cuf1
  79     CLEAR_SCREEN=clear CLEAR_EOL=el CLEAR_BOL=el1 CLEAR_EOS=ed BOLD=bold
  80     BLINK=blink DIM=dim REVERSE=rev UNDERLINE=smul NORMAL=sgr0
  81     HIDE_CURSOR=cinvis SHOW_CURSOR=cnorm""".split()
  82     _COLORS = """BLACK BLUE GREEN CYAN RED MAGENTA YELLOW WHITE""".split()
  83     _ANSICOLORS = "BLACK RED GREEN YELLOW BLUE MAGENTA CYAN WHITE".split()
  84
  85     def __init__(self, term_stream=sys.stdout):
  86         """
  87         Create a `TerminalController` and initialize its attributes
  88         with appropriate values for the current terminal.
  89         `term_stream` is the stream that will be used for terminal
  90         output; if this stream is not a tty, then the terminal is
  91         assumed to be a dumb terminal (i.e., have no capabilities).
  92         """
  93         # Curses isn't available on all platforms
  94         try: import curses
  95         except: return
  96
  97         # If the stream isn't a tty, then assume it has no capabilities.
  98         if not term_stream.isatty(): return
  99
 100         # Check the terminal type.  If we fail, then assume that the
 101         # terminal has no capabilities.
 102         try: curses.setupterm()
 103         except: return
 104
 105         # Look up numeric capabilities.
 106         self.COLS = curses.tigetnum('cols')
 107         self.LINES = curses.tigetnum('lines')
 108
 109         # Look up string capabilities.
 110         for capability in self._STRING_CAPABILITIES:
 111             (attrib, cap_name) = capability.split('=')
 112             setattr(self, attrib, self._tigetstr(cap_name) or '')
 113
 114         # Colors
 115         set_fg = self._tigetstr('setf')
 116         if set_fg:
 117             for i,color in zip(range(len(self._COLORS)), self._COLORS):
 118                 setattr(self, color, curses.tparm(set_fg, i) or '')
 119         set_fg_ansi = self._tigetstr('setaf')
 120         if set_fg_ansi:
 121             for i,color in zip(range(len(self._ANSICOLORS)), self._ANSICOLORS):
 122                 setattr(self, color, curses.tparm(set_fg_ansi, i) or '')
 123         set_bg = self._tigetstr('setb')
 124         if set_bg:
 125             for i,color in zip(range(len(self._COLORS)), self._COLORS):
 126                 setattr(self, 'BG_'+color, curses.tparm(set_bg, i) or '')
 127         set_bg_ansi = self._tigetstr('setab')
 128         if set_bg_ansi:
 129             for i,color in zip(range(len(self._ANSICOLORS)), self._ANSICOLORS):
 130                 setattr(self, 'BG_'+color, curses.tparm(set_bg_ansi, i) or '')
 131
 132     def _tigetstr(self, cap_name):
 133         # String capabilities can include "delays" of the form "$<2>".
 134         # For any modern terminal, we should be able to just ignore
 135         # these, so strip them out.
 136         import curses
 137         cap = curses.tigetstr(cap_name) or ''
 138         return re.sub(r'\$<\d+>[/*]?', '', cap)
 139
 140     def render(self, template):
 141         """
 142         Replace each $-substitutions in the given template string with
 143         the corresponding terminal control string (if it's defined) or
 144         '' (if it's not).
 145         """
 146         return re.sub(r'\$\$|\${\w+}', self._render_sub, template)
 147
 148     def _render_sub(self, match):
 149         s = match.group()
 150         if s == '$$': return s
 151         else: return getattr(self, s[2:-1])
 152
 153 #######################################################################
 154 # Example use case: progress bar
 155 #######################################################################
 156
 157 class SimpleProgressBar:
 158     """
 159     A simple progress bar which doesn't need any terminal support.
 160
 161     This prints out a progress bar like:
 162       'Header: 0 .. 10.. 20.. ...'
 163     """
 164
 165     def __init__(self, header):
 166         self.header = header
 167         self.atIndex = None
 168
 169     def update(self, percent, message):
 170         if self.atIndex is None:
 171             sys.stdout.write(self.header)
 172             self.atIndex = 0
 173
 174         next = int(percent*50)
 175         if next == self.atIndex:
 176             return
 177
 178         for i in range(self.atIndex, next):
 179             idx = i % 5
 180             if idx == 0:
 181                 sys.stdout.write('%-2d' % (i*2))
 182             elif idx == 1:
 183                 pass # Skip second char
 184             elif idx < 4:
 185                 sys.stdout.write('.')
 186             else:
 187                 sys.stdout.write(' ')
 188         sys.stdout.flush()
 189         self.atIndex = next
 190
 191     def clear(self):
 192         if self.atIndex is not None:
 193             sys.stdout.write('\n')
 194             sys.stdout.flush()
 195             self.atIndex = None
 196
 197 class ProgressBar:
 198     """
 199     A 3-line progress bar, which looks like::
 200
 201                                 Header
 202         20% [===========----------------------------------]
 203                            progress message
 204
 205     The progress bar is colored, if the terminal supports color
 206     output; and adjusts to the width of the terminal.
 207     """
 208     BAR = '%s${GREEN}[${BOLD}%s%s${NORMAL}${GREEN}]${NORMAL}%s\n'
 209     HEADER = '${BOLD}${CYAN}%s${NORMAL}\n\n'
 210
 211     def __init__(self, term, header, useETA=True):
 212         self.term = term
 213         if not (self.term.CLEAR_EOL and self.term.UP and self.term.BOL):
 214             raise ValueError("Terminal isn't capable enough -- you "
 215                              "should use a simpler progress dispaly.")
 216         self.width = self.term.COLS or 75
 217         self.bar = term.render(self.BAR)
 218         self.header = self.term.render(self.HEADER % header.center(self.width))
 219         self.cleared = 1 #: true if we haven't drawn the bar yet.
 220         self.useETA = useETA
 221         if self.useETA:
 222             self.startTime = time.time()
 223         self.update(0, '')
 224
 225     def update(self, percent, message):
 226         if self.cleared:
 227             sys.stdout.write(self.header)
 228             self.cleared = 0
 229         prefix = '%3d%% ' % (percent*100,)
 230         suffix = ''
 231         if self.useETA:
 232             elapsed = time.time() - self.startTime
 233             if percent > .0001 and elapsed > 1:
 234                 total = elapsed / percent
 235                 eta = int(total - elapsed)
 236                 h = eta//3600.
 237                 m = (eta//60) % 60
 238                 s = eta % 60
 239                 suffix = ' ETA: %02d:%02d:%02d'%(h,m,s)
 240         barWidth = self.width - len(prefix) - len(suffix) - 2
 241         n = int(barWidth*percent)
 242         if len(message) < self.width:
 243             message = message + ' '*(self.width - len(message))
 244         else:
 245             message = '... ' + message[-(self.width-4):]
 246         sys.stdout.write(
 247             self.term.BOL + self.term.UP + self.term.CLEAR_EOL +
 248             (self.bar % (prefix, '='*n, '-'*(barWidth-n), suffix)) +
 249             self.term.CLEAR_EOL + message)
 250
 251     def clear(self):
 252         if not self.cleared:
 253             sys.stdout.write(self.term.BOL + self.term.CLEAR_EOL +
 254                              self.term.UP + self.term.CLEAR_EOL +
 255                              self.term.UP + self.term.CLEAR_EOL)
 256             self.cleared = 1
 257
 258 def test():
 259     import time
 260     tc = TerminalController()
 261     p = ProgressBar(tc, 'Tests')
 262     for i in range(101):
 263         p.update(i/100., str(i))
 264         time.sleep(.3)
 265
 266 if __name__=='__main__':
 267     test()