Home | History | Annotate | Download | only in lit
      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