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