1#!/usr/bin/env python 2 3# Source: http://code.activestate.com/recipes/475116/, with 4# modifications by Daniel Dunbar. 5 6import sys, re, time 7 8def to_bytes(str): 9 # Encode to UTF-8 to get binary data. 10 return str.encode('utf-8') 11 12class 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 170class 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, interrupted): 205 if self.atIndex is not None and not interrupted: 206 sys.stdout.write('\n') 207 sys.stdout.flush() 208 self.atIndex = None 209 210class 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${%s}[${BOLD}%s%s${NORMAL}${%s}]${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.barColor = 'GREEN' 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 = 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 bc = self.barColor 268 bar = self.BAR % (prefix, bc, '='*n, '-'*(barWidth-n), bc, suffix) 269 bar = self.term.render(bar) 270 sys.stdout.write( 271 self.BOL + self.term.UP + self.term.CLEAR_EOL + 272 bar + 273 self.XNL + 274 self.term.CLEAR_EOL + message) 275 if not self.term.XN: 276 sys.stdout.flush() 277 278 def clear(self, interrupted): 279 if not self.cleared: 280 sys.stdout.write(self.BOL + self.term.CLEAR_EOL + 281 self.term.UP + self.term.CLEAR_EOL + 282 self.term.UP + self.term.CLEAR_EOL) 283 if interrupted: # ^C creates extra line. Gobble it up! 284 sys.stdout.write(self.term.UP + self.term.CLEAR_EOL) 285 sys.stdout.write('^C') 286 sys.stdout.flush() 287 self.cleared = 1 288 289def test(): 290 tc = TerminalController() 291 p = ProgressBar(tc, 'Tests') 292 for i in range(101): 293 p.update(i/100., str(i)) 294 time.sleep(.3) 295 296if __name__=='__main__': 297 test() 298