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() 281
