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
8class 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
158class 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
198class 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
271def 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
279if __name__=='__main__':
280    test()
281