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