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 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         self.XN = curses.tigetflag('xenl')
 109
 110         # Look up string capabilities.
 111         for capability in self._STRING_CAPABILITIES:
 112             (attrib, cap_name) = capability.split('=')
 113             setattr(self, attrib, self._tigetstr(cap_name) or '')
 114
 115         # Colors
 116         set_fg = self._tigetstr('setf')
 117         if set_fg:
 118             for i,color in zip(range(len(self._COLORS)), self._COLORS):
 119                 setattr(self, color, curses.tparm(set_fg, i) or '')
 120         set_fg_ansi = self._tigetstr('setaf')
 121         if set_fg_ansi:
 122             for i,color in zip(range(len(self._ANSICOLORS)), self._ANSICOLORS):
 123                 setattr(self, color, curses.tparm(set_fg_ansi, i) or '')
 124         set_bg = self._tigetstr('setb')
 125         if set_bg:
 126             for i,color in zip(range(len(self._COLORS)), self._COLORS):
 127                 setattr(self, 'BG_'+color, curses.tparm(set_bg, i) or '')
 128         set_bg_ansi = self._tigetstr('setab')
 129         if set_bg_ansi:
 130             for i,color in zip(range(len(self._ANSICOLORS)), self._ANSICOLORS):
 131                 setattr(self, 'BG_'+color, curses.tparm(set_bg_ansi, i) or '')
 132
 133     def _tigetstr(self, cap_name):
 134         # String capabilities can include "delays" of the form "$<2>".
 135         # For any modern terminal, we should be able to just ignore
 136         # these, so strip them out.
 137         import curses
 138         cap = curses.tigetstr(cap_name) or ''
 139         return re.sub(r'\$<\d+>[/*]?', '', cap)
 140
 141     def render(self, template):
 142         """
 143         Replace each $-substitutions in the given template string with
 144         the corresponding terminal control string (if it's defined) or
 145         '' (if it's not).
 146         """
 147         return re.sub(r'\$\$|\${\w+}', self._render_sub, template)
 148
 149     def _render_sub(self, match):
 150         s = match.group()
 151         if s == '$$': return s
 152         else: return getattr(self, s[2:-1])
 153
 154 #######################################################################
 155 # Example use case: progress bar
 156 #######################################################################
 157
 158 class SimpleProgressBar:
 159     """
 160     A simple progress bar which doesn't need any terminal support.
 161
 162     This prints out a progress bar like:
 163       'Header: 0 .. 10.. 20.. ...'
 164     """
 165
 166     def __init__(self, header):
 167         self.header = header
 168         self.atIndex = None
 169
 170     def update(self, percent, message):
 171         if self.atIndex is None:
 172             sys.stdout.write(self.header)
 173             self.atIndex = 0
 174
 175         next = int(percent*50)
 176         if next == self.atIndex:
 177             return
 178
 179         for i in range(self.atIndex, next):
 180             idx = i % 5
 181             if idx == 0:
 182                 sys.stdout.write('%-2d' % (i*2))
 183             elif idx == 1:
 184                 pass # Skip second char
 185             elif idx < 4:
 186                 sys.stdout.write('.')
 187             else:
 188                 sys.stdout.write(' ')
 189         sys.stdout.flush()
 190         self.atIndex = next
 191
 192     def clear(self):
 193         if self.atIndex is not None:
 194             sys.stdout.write('\n')
 195             sys.stdout.flush()
 196             self.atIndex = None
 197
 198 class ProgressBar:
 199     """
 200     A 3-line progress bar, which looks like::
 201
 202                                 Header
 203         20% [===========----------------------------------]
 204                            progress message
 205
 206     The progress bar is colored, if the terminal supports color
 207     output; and adjusts to the width of the terminal.
 208     """
 209     BAR = '%s${GREEN}[${BOLD}%s%s${NORMAL}${GREEN}]${NORMAL}%s'
 210     HEADER = '${BOLD}${CYAN}%s${NORMAL}\n\n'
 211
 212     def __init__(self, term, header, useETA=True):
 213         self.term = term
 214         if not (self.term.CLEAR_EOL and self.term.UP and self.term.BOL):
 215             raise ValueError("Terminal isn't capable enough -- you "
 216                              "should use a simpler progress dispaly.")
 217         self.BOL = self.term.BOL # BoL from col#79
 218         self.XNL = "\n" # Newline from col#79
 219         if self.term.COLS:
 220             self.width = self.term.COLS
 221             if not self.term.XN:
 222                 self.BOL = self.term.UP + self.term.BOL
 223                 self.XNL = "" # Cursor must be fed to the next line
 224         else:
 225             self.width = 75
 226         self.bar = term.render(self.BAR)
 227         self.header = self.term.render(self.HEADER % header.center(self.width))
 228         self.cleared = 1 #: true if we haven't drawn the bar yet.
 229         self.useETA = useETA
 230         if self.useETA:
 231             self.startTime = time.time()
 232         self.update(0, '')
 233
 234     def update(self, percent, message):
 235         if self.cleared:
 236             sys.stdout.write(self.header)
 237             self.cleared = 0
 238         prefix = '%3d%% ' % (percent*100,)
 239         suffix = ''
 240         if self.useETA:
 241             elapsed = time.time() - self.startTime
 242             if percent > .0001 and elapsed > 1:
 243                 total = elapsed / percent
 244                 eta = int(total - elapsed)
 245                 h = eta//3600.
 246                 m = (eta//60) % 60
 247                 s = eta % 60
 248                 suffix = ' ETA: %02d:%02d:%02d'%(h,m,s)
 249         barWidth = self.width - len(prefix) - len(suffix) - 2
 250         n = int(barWidth*percent)
 251         if len(message) < self.width:
 252             message = message + ' '*(self.width - len(message))
 253         else:
 254             message = '... ' + message[-(self.width-4):]
 255         sys.stdout.write(
 256             self.BOL + self.term.UP + self.term.CLEAR_EOL +
 257             (self.bar % (prefix, '='*n, '-'*(barWidth-n), suffix)) +
 258             self.XNL +
 259             self.term.CLEAR_EOL + message)
 260         if not self.term.XN:
 261             sys.stdout.flush()
 262
 263     def clear(self):
 264         if not self.cleared:
 265             sys.stdout.write(self.BOL + self.term.CLEAR_EOL +
 266                              self.term.UP + self.term.CLEAR_EOL +
 267                              self.term.UP + self.term.CLEAR_EOL)
 268             sys.stdout.flush()
 269             self.cleared = 1
 270
 271 def test():
 272     import time
 273     tc = TerminalController()
 274     p = ProgressBar(tc, 'Tests')
 275     for i in range(101):
 276         p.update(i/100., str(i))
 277         time.sleep(.3)
 278
 279 if __name__=='__main__':
 280     test()