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