109467b48Spatrick#!/usr/bin/env python 209467b48Spatrick 309467b48Spatrick# Source: http://code.activestate.com/recipes/475116/, with 409467b48Spatrick# modifications by Daniel Dunbar. 509467b48Spatrick 609467b48Spatrickimport sys, re, time 709467b48Spatrick 809467b48Spatrickdef to_bytes(str): 909467b48Spatrick # Encode to UTF-8 to get binary data. 1009467b48Spatrick return str.encode('utf-8') 1109467b48Spatrick 1209467b48Spatrickclass TerminalController: 1309467b48Spatrick """ 1409467b48Spatrick A class that can be used to portably generate formatted output to 1509467b48Spatrick a terminal. 1609467b48Spatrick 1709467b48Spatrick `TerminalController` defines a set of instance variables whose 1809467b48Spatrick values are initialized to the control sequence necessary to 1909467b48Spatrick perform a given action. These can be simply included in normal 2009467b48Spatrick output to the terminal: 2109467b48Spatrick 2209467b48Spatrick >>> term = TerminalController() 2309467b48Spatrick >>> print('This is '+term.GREEN+'green'+term.NORMAL) 2409467b48Spatrick 2509467b48Spatrick Alternatively, the `render()` method can used, which replaces 2609467b48Spatrick '${action}' with the string required to perform 'action': 2709467b48Spatrick 2809467b48Spatrick >>> term = TerminalController() 2909467b48Spatrick >>> print(term.render('This is ${GREEN}green${NORMAL}')) 3009467b48Spatrick 3109467b48Spatrick If the terminal doesn't support a given action, then the value of 3209467b48Spatrick the corresponding instance variable will be set to ''. As a 3309467b48Spatrick result, the above code will still work on terminals that do not 3409467b48Spatrick support color, except that their output will not be colored. 3509467b48Spatrick Also, this means that you can test whether the terminal supports a 3609467b48Spatrick given action by simply testing the truth value of the 3709467b48Spatrick corresponding instance variable: 3809467b48Spatrick 3909467b48Spatrick >>> term = TerminalController() 4009467b48Spatrick >>> if term.CLEAR_SCREEN: 4109467b48Spatrick ... print('This terminal supports clearning the screen.') 4209467b48Spatrick 4309467b48Spatrick Finally, if the width and height of the terminal are known, then 4409467b48Spatrick they will be stored in the `COLS` and `LINES` attributes. 4509467b48Spatrick """ 4609467b48Spatrick # Cursor movement: 4709467b48Spatrick BOL = '' #: Move the cursor to the beginning of the line 4809467b48Spatrick UP = '' #: Move the cursor up one line 4909467b48Spatrick DOWN = '' #: Move the cursor down one line 5009467b48Spatrick LEFT = '' #: Move the cursor left one char 5109467b48Spatrick RIGHT = '' #: Move the cursor right one char 5209467b48Spatrick 5309467b48Spatrick # Deletion: 5409467b48Spatrick CLEAR_SCREEN = '' #: Clear the screen and move to home position 5509467b48Spatrick CLEAR_EOL = '' #: Clear to the end of the line. 5609467b48Spatrick CLEAR_BOL = '' #: Clear to the beginning of the line. 5709467b48Spatrick CLEAR_EOS = '' #: Clear to the end of the screen 5809467b48Spatrick 5909467b48Spatrick # Output modes: 6009467b48Spatrick BOLD = '' #: Turn on bold mode 6109467b48Spatrick BLINK = '' #: Turn on blink mode 6209467b48Spatrick DIM = '' #: Turn on half-bright mode 6309467b48Spatrick REVERSE = '' #: Turn on reverse-video mode 6409467b48Spatrick NORMAL = '' #: Turn off all modes 6509467b48Spatrick 6609467b48Spatrick # Cursor display: 6709467b48Spatrick HIDE_CURSOR = '' #: Make the cursor invisible 6809467b48Spatrick SHOW_CURSOR = '' #: Make the cursor visible 6909467b48Spatrick 7009467b48Spatrick # Terminal size: 7109467b48Spatrick COLS = None #: Width of the terminal (None for unknown) 7209467b48Spatrick LINES = None #: Height of the terminal (None for unknown) 7309467b48Spatrick 7409467b48Spatrick # Foreground colors: 7509467b48Spatrick BLACK = BLUE = GREEN = CYAN = RED = MAGENTA = YELLOW = WHITE = '' 7609467b48Spatrick 7709467b48Spatrick # Background colors: 7809467b48Spatrick BG_BLACK = BG_BLUE = BG_GREEN = BG_CYAN = '' 7909467b48Spatrick BG_RED = BG_MAGENTA = BG_YELLOW = BG_WHITE = '' 8009467b48Spatrick 8109467b48Spatrick _STRING_CAPABILITIES = """ 8209467b48Spatrick BOL=cr UP=cuu1 DOWN=cud1 LEFT=cub1 RIGHT=cuf1 8309467b48Spatrick CLEAR_SCREEN=clear CLEAR_EOL=el CLEAR_BOL=el1 CLEAR_EOS=ed BOLD=bold 8409467b48Spatrick BLINK=blink DIM=dim REVERSE=rev UNDERLINE=smul NORMAL=sgr0 8509467b48Spatrick HIDE_CURSOR=cinvis SHOW_CURSOR=cnorm""".split() 8609467b48Spatrick _COLORS = """BLACK BLUE GREEN CYAN RED MAGENTA YELLOW WHITE""".split() 8709467b48Spatrick _ANSICOLORS = "BLACK RED GREEN YELLOW BLUE MAGENTA CYAN WHITE".split() 8809467b48Spatrick 8909467b48Spatrick def __init__(self, term_stream=sys.stdout): 9009467b48Spatrick """ 9109467b48Spatrick Create a `TerminalController` and initialize its attributes 9209467b48Spatrick with appropriate values for the current terminal. 9309467b48Spatrick `term_stream` is the stream that will be used for terminal 9409467b48Spatrick output; if this stream is not a tty, then the terminal is 9509467b48Spatrick assumed to be a dumb terminal (i.e., have no capabilities). 9609467b48Spatrick """ 9709467b48Spatrick # Curses isn't available on all platforms 9809467b48Spatrick try: import curses 9909467b48Spatrick except: return 10009467b48Spatrick 10109467b48Spatrick # If the stream isn't a tty, then assume it has no capabilities. 10209467b48Spatrick if not term_stream.isatty(): return 10309467b48Spatrick 10409467b48Spatrick # Check the terminal type. If we fail, then assume that the 10509467b48Spatrick # terminal has no capabilities. 10609467b48Spatrick try: curses.setupterm() 10709467b48Spatrick except: return 10809467b48Spatrick 10909467b48Spatrick # Look up numeric capabilities. 11009467b48Spatrick self.COLS = curses.tigetnum('cols') 11109467b48Spatrick self.LINES = curses.tigetnum('lines') 11209467b48Spatrick self.XN = curses.tigetflag('xenl') 11309467b48Spatrick 11409467b48Spatrick # Look up string capabilities. 11509467b48Spatrick for capability in self._STRING_CAPABILITIES: 11609467b48Spatrick (attrib, cap_name) = capability.split('=') 11709467b48Spatrick setattr(self, attrib, self._tigetstr(cap_name) or '') 11809467b48Spatrick 11909467b48Spatrick # Colors 12009467b48Spatrick set_fg = self._tigetstr('setf') 12109467b48Spatrick if set_fg: 12209467b48Spatrick for i,color in zip(range(len(self._COLORS)), self._COLORS): 12309467b48Spatrick setattr(self, color, self._tparm(set_fg, i)) 12409467b48Spatrick set_fg_ansi = self._tigetstr('setaf') 12509467b48Spatrick if set_fg_ansi: 12609467b48Spatrick for i,color in zip(range(len(self._ANSICOLORS)), self._ANSICOLORS): 12709467b48Spatrick setattr(self, color, self._tparm(set_fg_ansi, i)) 12809467b48Spatrick set_bg = self._tigetstr('setb') 12909467b48Spatrick if set_bg: 13009467b48Spatrick for i,color in zip(range(len(self._COLORS)), self._COLORS): 13109467b48Spatrick setattr(self, 'BG_'+color, self._tparm(set_bg, i)) 13209467b48Spatrick set_bg_ansi = self._tigetstr('setab') 13309467b48Spatrick if set_bg_ansi: 13409467b48Spatrick for i,color in zip(range(len(self._ANSICOLORS)), self._ANSICOLORS): 13509467b48Spatrick setattr(self, 'BG_'+color, self._tparm(set_bg_ansi, i)) 13609467b48Spatrick 13709467b48Spatrick def _tparm(self, arg, index): 13809467b48Spatrick import curses 13909467b48Spatrick return curses.tparm(to_bytes(arg), index).decode('utf-8') or '' 14009467b48Spatrick 14109467b48Spatrick def _tigetstr(self, cap_name): 14209467b48Spatrick # String capabilities can include "delays" of the form "$<2>". 14309467b48Spatrick # For any modern terminal, we should be able to just ignore 14409467b48Spatrick # these, so strip them out. 14509467b48Spatrick import curses 14609467b48Spatrick cap = curses.tigetstr(cap_name) 14709467b48Spatrick if cap is None: 14809467b48Spatrick cap = '' 14909467b48Spatrick else: 15009467b48Spatrick cap = cap.decode('utf-8') 15109467b48Spatrick return re.sub(r'\$<\d+>[/*]?', '', cap) 15209467b48Spatrick 15309467b48Spatrick def render(self, template): 15409467b48Spatrick """ 15509467b48Spatrick Replace each $-substitutions in the given template string with 15609467b48Spatrick the corresponding terminal control string (if it's defined) or 15709467b48Spatrick '' (if it's not). 15809467b48Spatrick """ 15909467b48Spatrick return re.sub(r'\$\$|\${\w+}', self._render_sub, template) 16009467b48Spatrick 16109467b48Spatrick def _render_sub(self, match): 16209467b48Spatrick s = match.group() 16309467b48Spatrick if s == '$$': return s 16409467b48Spatrick else: return getattr(self, s[2:-1]) 16509467b48Spatrick 16609467b48Spatrick####################################################################### 16709467b48Spatrick# Example use case: progress bar 16809467b48Spatrick####################################################################### 16909467b48Spatrick 17009467b48Spatrickclass SimpleProgressBar: 17109467b48Spatrick """ 17209467b48Spatrick A simple progress bar which doesn't need any terminal support. 17309467b48Spatrick 17409467b48Spatrick This prints out a progress bar like: 17509467b48Spatrick 'Header: 0.. 10.. 20.. ...' 17609467b48Spatrick """ 17709467b48Spatrick 17809467b48Spatrick def __init__(self, header): 17909467b48Spatrick self.header = header 18009467b48Spatrick self.atIndex = None 18109467b48Spatrick 18209467b48Spatrick def update(self, percent, message): 18309467b48Spatrick if self.atIndex is None: 18409467b48Spatrick sys.stdout.write(self.header) 18509467b48Spatrick self.atIndex = 0 18609467b48Spatrick 18709467b48Spatrick next = int(percent*50) 18809467b48Spatrick if next == self.atIndex: 18909467b48Spatrick return 19009467b48Spatrick 19109467b48Spatrick for i in range(self.atIndex, next): 19209467b48Spatrick idx = i % 5 19309467b48Spatrick if idx == 0: 19409467b48Spatrick sys.stdout.write('%2d' % (i*2)) 19509467b48Spatrick elif idx == 1: 19609467b48Spatrick pass # Skip second char 19709467b48Spatrick elif idx < 4: 19809467b48Spatrick sys.stdout.write('.') 19909467b48Spatrick else: 20009467b48Spatrick sys.stdout.write(' ') 20109467b48Spatrick sys.stdout.flush() 20209467b48Spatrick self.atIndex = next 20309467b48Spatrick 20409467b48Spatrick def clear(self, interrupted): 20509467b48Spatrick if self.atIndex is not None and not interrupted: 20609467b48Spatrick sys.stdout.write('\n') 20709467b48Spatrick sys.stdout.flush() 20809467b48Spatrick self.atIndex = None 20909467b48Spatrick 21009467b48Spatrickclass ProgressBar: 21109467b48Spatrick """ 21209467b48Spatrick A 3-line progress bar, which looks like:: 21309467b48Spatrick 21409467b48Spatrick Header 21509467b48Spatrick 20% [===========----------------------------------] 21609467b48Spatrick progress message 21709467b48Spatrick 21809467b48Spatrick The progress bar is colored, if the terminal supports color 21909467b48Spatrick output; and adjusts to the width of the terminal. 22009467b48Spatrick """ 22109467b48Spatrick BAR = '%s${%s}[${BOLD}%s%s${NORMAL}${%s}]${NORMAL}%s' 22209467b48Spatrick HEADER = '${BOLD}${CYAN}%s${NORMAL}\n\n' 22309467b48Spatrick 22409467b48Spatrick def __init__(self, term, header, useETA=True): 22509467b48Spatrick self.term = term 22609467b48Spatrick if not (self.term.CLEAR_EOL and self.term.UP and self.term.BOL): 22709467b48Spatrick raise ValueError("Terminal isn't capable enough -- you " 22809467b48Spatrick "should use a simpler progress dispaly.") 22909467b48Spatrick self.BOL = self.term.BOL # BoL from col#79 23009467b48Spatrick self.XNL = "\n" # Newline from col#79 23109467b48Spatrick if self.term.COLS: 23209467b48Spatrick self.width = self.term.COLS 23309467b48Spatrick if not self.term.XN: 23409467b48Spatrick self.BOL = self.term.UP + self.term.BOL 23509467b48Spatrick self.XNL = "" # Cursor must be fed to the next line 23609467b48Spatrick else: 23709467b48Spatrick self.width = 75 23809467b48Spatrick self.barColor = 'GREEN' 23909467b48Spatrick self.header = self.term.render(self.HEADER % header.center(self.width)) 24009467b48Spatrick self.cleared = 1 #: true if we haven't drawn the bar yet. 24109467b48Spatrick self.useETA = useETA 24209467b48Spatrick if self.useETA: 24309467b48Spatrick self.startTime = time.time() 24409467b48Spatrick # self.update(0, '') 24509467b48Spatrick 24609467b48Spatrick def update(self, percent, message): 24709467b48Spatrick if self.cleared: 24809467b48Spatrick sys.stdout.write(self.header) 24909467b48Spatrick self.cleared = 0 25009467b48Spatrick prefix = '%3d%% ' % (percent*100,) 25109467b48Spatrick suffix = '' 25209467b48Spatrick if self.useETA: 25309467b48Spatrick elapsed = time.time() - self.startTime 25409467b48Spatrick if percent > .0001 and elapsed > 1: 25509467b48Spatrick total = elapsed / percent 256*73471bf0Spatrick eta = total - elapsed 25709467b48Spatrick h = eta//3600. 25809467b48Spatrick m = (eta//60) % 60 25909467b48Spatrick s = eta % 60 26009467b48Spatrick suffix = ' ETA: %02d:%02d:%02d'%(h,m,s) 26109467b48Spatrick barWidth = self.width - len(prefix) - len(suffix) - 2 26209467b48Spatrick n = int(barWidth*percent) 26309467b48Spatrick if len(message) < self.width: 26409467b48Spatrick message = message + ' '*(self.width - len(message)) 26509467b48Spatrick else: 26609467b48Spatrick message = '... ' + message[-(self.width-4):] 26709467b48Spatrick bc = self.barColor 26809467b48Spatrick bar = self.BAR % (prefix, bc, '='*n, '-'*(barWidth-n), bc, suffix) 26909467b48Spatrick bar = self.term.render(bar) 27009467b48Spatrick sys.stdout.write( 27109467b48Spatrick self.BOL + self.term.UP + self.term.CLEAR_EOL + 27209467b48Spatrick bar + 27309467b48Spatrick self.XNL + 27409467b48Spatrick self.term.CLEAR_EOL + message) 27509467b48Spatrick if not self.term.XN: 27609467b48Spatrick sys.stdout.flush() 27709467b48Spatrick 27809467b48Spatrick def clear(self, interrupted): 27909467b48Spatrick if not self.cleared: 28009467b48Spatrick sys.stdout.write(self.BOL + self.term.CLEAR_EOL + 28109467b48Spatrick self.term.UP + self.term.CLEAR_EOL + 28209467b48Spatrick self.term.UP + self.term.CLEAR_EOL) 28309467b48Spatrick if interrupted: # ^C creates extra line. Gobble it up! 28409467b48Spatrick sys.stdout.write(self.term.UP + self.term.CLEAR_EOL) 28509467b48Spatrick sys.stdout.write('^C') 28609467b48Spatrick sys.stdout.flush() 28709467b48Spatrick self.cleared = 1 28809467b48Spatrick 28909467b48Spatrickdef test(): 29009467b48Spatrick tc = TerminalController() 29109467b48Spatrick p = ProgressBar(tc, 'Tests') 29209467b48Spatrick for i in range(101): 29309467b48Spatrick p.update(i/100., str(i)) 29409467b48Spatrick time.sleep(.3) 29509467b48Spatrick 29609467b48Spatrickif __name__=='__main__': 29709467b48Spatrick test() 298