From d00a37e66acb95bac4970e9b5fb719e75c488cb1 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Tue, 20 Nov 2012 16:17:44 +0530 Subject: [PATCH] Enable colored output for the CLI on windows as well --- src/calibre/constants.py | 10 +- src/calibre/customize/builtins.py | 2 +- src/calibre/devices/cli.py | 20 +-- src/calibre/library/cli.py | 16 +- src/calibre/utils/config.py | 37 +++-- src/calibre/utils/logging.py | 17 +-- src/calibre/utils/terminal.py | 239 ++++++++++++++++++++++++++++++ src/calibre/utils/terminfo.py | 215 --------------------------- 8 files changed, 290 insertions(+), 266 deletions(-) create mode 100644 src/calibre/utils/terminal.py delete mode 100644 src/calibre/utils/terminfo.py diff --git a/src/calibre/constants.py b/src/calibre/constants.py index b6fa0b42d7..d15825d4b9 100644 --- a/src/calibre/constants.py +++ b/src/calibre/constants.py @@ -14,14 +14,6 @@ Various run time constants. import sys, locale, codecs, os, importlib, collections -_tc = None -def terminal_controller(): - global _tc - if _tc is None: - from calibre.utils.terminfo import TerminalController - _tc = TerminalController(sys.stdout) - return _tc - _plat = sys.platform.lower() iswindows = 'win32' in _plat or 'win64' in _plat isosx = 'darwin' in _plat @@ -37,6 +29,8 @@ isportable = os.environ.get('CALIBRE_PORTABLE_BUILD', None) is not None ispy3 = sys.version_info.major > 2 isxp = iswindows and sys.getwindowsversion().major < 6 isworker = os.environ.has_key('CALIBRE_WORKER') or os.environ.has_key('CALIBRE_SIMPLE_WORKER') +if isworker: + os.environ.pop('CALIBRE_FORCE_ANSI', None) try: preferred_encoding = locale.getpreferredencoding() diff --git a/src/calibre/customize/builtins.py b/src/calibre/customize/builtins.py index 693a9c4a6a..1af944cf4f 100644 --- a/src/calibre/customize/builtins.py +++ b/src/calibre/customize/builtins.py @@ -1716,7 +1716,7 @@ if __name__ == '__main__': ret = 0 for x in ('lxml', 'calibre.ebooks.BeautifulSoup', 'uuid', - 'calibre.utils.terminfo', 'calibre.utils.magick', 'PIL', 'Image', + 'calibre.utils.terminal', 'calibre.utils.magick', 'PIL', 'Image', 'sqlite3', 'mechanize', 'httplib', 'xml'): if x in sys.modules: ret = 1 diff --git a/src/calibre/devices/cli.py b/src/calibre/devices/cli.py index 0d5c7f2f19..e1a5375b6b 100755 --- a/src/calibre/devices/cli.py +++ b/src/calibre/devices/cli.py @@ -11,7 +11,6 @@ from optparse import OptionParser from calibre import __version__, __appname__, human_readable from calibre.devices.errors import PathError -from calibre.utils.terminfo import TerminalController from calibre.devices.errors import ArgumentError, DeviceError, DeviceLocked from calibre.customize.ui import device_plugins from calibre.devices.scanner import DeviceScanner @@ -20,8 +19,7 @@ from calibre.utils.config import device_prefs MINIMUM_COL_WIDTH = 12 #: Minimum width of columns in ls output class FileFormatter(object): - def __init__(self, file, term): - self.term = term + def __init__(self, file): self.is_dir = file.is_dir self.is_readonly = file.is_readonly self.size = file.size @@ -94,7 +92,7 @@ def info(dev): print "Software version:", info[2] print "Mime type: ", info[3] -def ls(dev, path, term, recurse=False, color=False, human_readable_size=False, ll=False, cols=0): +def ls(dev, path, recurse=False, human_readable_size=False, ll=False, cols=0): def col_split(l, cols): # split list l into columns rows = len(l) / cols if len(l) % cols: @@ -126,14 +124,13 @@ def ls(dev, path, term, recurse=False, color=False, human_readable_size=False, l for file in files: size = len(str(file.size)) if human_readable_size: - file = FileFormatter(file, term) + file = FileFormatter(file) size = len(file.human_readable_size) if size > maxlen: maxlen = size for file in files: - file = FileFormatter(file, term) + file = FileFormatter(file) name = file.name if ll else file.isdir_name lsoutput.append(name) - if color: name = file.name_in_color lscoloutput.append(name) if ll: size = str(file.size) @@ -173,10 +170,8 @@ def shutdown_plugins(): pass def main(): - term = TerminalController() - cols = term.COLS - if not cols: # On windows terminal width is unknown - cols = 80 + from calibre.utils.terminal import geometry + cols = geometry()[0] parser = OptionParser(usage="usage: %prog [options] command args\n\ncommand "+ "is one of: info, books, df, ls, cp, mkdir, touch, cat, rm, eject, test_file\n\n"+ @@ -260,7 +255,6 @@ def main(): dev.mkdir(args[0]) elif command == "ls": parser = OptionParser(usage="usage: %prog ls [options] path\nList files on the device\n\npath must begin with / or card:/") - parser.add_option("--color", help="show ls output in color", dest="color", action="store_true", default=False) parser.add_option("-l", help="In addition to the name of each file, print the file type, permissions, and timestamp (the modification time, in the local timezone). Times are local.", dest="ll", action="store_true", default=False) parser.add_option("-R", help="Recursively list subdirectories encountered. /dev and /proc are omitted", dest="recurse", action="store_true", default=False) parser.remove_option("-h") @@ -269,7 +263,7 @@ def main(): if len(args) != 1: parser.print_help() return 1 - print ls(dev, args[0], term, color=options.color, recurse=options.recurse, ll=options.ll, human_readable_size=options.hrs, cols=cols), + print ls(dev, args[0], recurse=options.recurse, ll=options.ll, human_readable_size=options.hrs, cols=cols), elif command == "info": info(dev) elif command == "cp": diff --git a/src/calibre/library/cli.py b/src/calibre/library/cli.py index 3993244215..676d3cd1bc 100644 --- a/src/calibre/library/cli.py +++ b/src/calibre/library/cli.py @@ -65,8 +65,7 @@ def get_db(dbpath, options): def do_list(db, fields, afields, sort_by, ascending, search_text, line_width, separator, prefix, subtitle='Books in the calibre database'): - from calibre.constants import terminal_controller as tc - terminal_controller = tc() + from calibre.utils.terminal import ColoredStream, geometry if sort_by: db.sort(sort_by, ascending) if search_text: @@ -101,7 +100,7 @@ def do_list(db, fields, afields, sort_by, ascending, search_text, line_width, se for j, field in enumerate(fields): widths[j] = max(widths[j], len(unicode(i[field]))) - screen_width = terminal_controller.COLS if line_width < 0 else line_width + screen_width = geometry()[0] if line_width < 0 else line_width if not screen_width: screen_width = 80 field_width = screen_width//len(fields) @@ -120,7 +119,8 @@ def do_list(db, fields, afields, sort_by, ascending, search_text, line_width, se widths = list(base_widths) titles = map(lambda x, y: '%-*s%s'%(x-len(separator), y, separator), widths, title_fields) - print terminal_controller.GREEN + ''.join(titles)+terminal_controller.NORMAL + with ColoredStream(sys.stdout, fg='green'): + print ''.join(titles) wrappers = map(lambda x: TextWrapper(x-1), widths) o = cStringIO.StringIO() @@ -1288,8 +1288,7 @@ def command_list_categories(args, dbpath): fields = ['category', 'tag_name', 'count', 'rating'] def do_list(): - from calibre.constants import terminal_controller as tc - terminal_controller = tc() + from calibre.utils.terminal import geometry, ColoredStream separator = ' ' widths = list(map(lambda x : 0, fields)) @@ -1297,7 +1296,7 @@ def command_list_categories(args, dbpath): for j, field in enumerate(fields): widths[j] = max(widths[j], max(len(field), len(unicode(i[field])))) - screen_width = terminal_controller.COLS if opts.width < 0 else opts.width + screen_width = geometry()[0] if not screen_width: screen_width = 80 field_width = screen_width//len(fields) @@ -1316,7 +1315,8 @@ def command_list_categories(args, dbpath): widths = list(base_widths) titles = map(lambda x, y: '%-*s%s'%(x-len(separator), y, separator), widths, fields) - print terminal_controller.GREEN + ''.join(titles)+terminal_controller.NORMAL + with ColoredStream(sys.stdout, fg='green'): + print ''.join(titles) wrappers = map(lambda x: TextWrapper(x-1), widths) o = cStringIO.StringIO() diff --git a/src/calibre/utils/config.py b/src/calibre/utils/config.py index e79affb2f4..2f562ae90e 100644 --- a/src/calibre/utils/config.py +++ b/src/calibre/utils/config.py @@ -12,7 +12,7 @@ from optparse import OptionParser as _OptionParser, OptionGroup from optparse import IndentedHelpFormatter from calibre.constants import (config_dir, CONFIG_DIR_MODE, __appname__, - __version__, __author__, terminal_controller) + __version__, __author__) from calibre.utils.lock import ExclusiveFile from calibre.utils.config_base import (make_config_dir, Option, OptionValues, OptionSet, ConfigInterface, Config, prefs, StringConfig, ConfigProxy, @@ -30,28 +30,28 @@ def check_config_write_access(): class CustomHelpFormatter(IndentedHelpFormatter): def format_usage(self, usage): - tc = terminal_controller() - return "%s%s%s: %s\n" % (tc.BLUE, _('Usage'), tc.NORMAL, usage) + from calibre.utils.terminal import colored + return colored(_('Usage'), fg='blue', bold=True) + ': ' + usage def format_heading(self, heading): - tc = terminal_controller() - return "%*s%s%s%s:\n" % (self.current_indent, tc.BLUE, - "", heading, tc.NORMAL) + from calibre.utils.terminal import colored + return "%*s%s:\n" % (self.current_indent, '', + colored(heading, fg='blue')) def format_option(self, option): import textwrap - tc = terminal_controller() + from calibre.utils.terminal import colored result = [] opts = self.option_strings[option] opt_width = self.help_position - self.current_indent - 2 if len(opts) > opt_width: opts = "%*s%s\n" % (self.current_indent, "", - tc.GREEN+opts+tc.NORMAL) + colored(opts, fg='green')) indent_first = self.help_position else: # start help on same line as opts opts = "%*s%-*s " % (self.current_indent, "", opt_width + - len(tc.GREEN + tc.NORMAL), tc.GREEN + opts + tc.NORMAL) + len(colored('', fg='green')), colored(opts, fg='green')) indent_first = 0 result.append(opts) if option.help: @@ -78,11 +78,11 @@ class OptionParser(_OptionParser): conflict_handler='resolve', **kwds): import textwrap - tc = terminal_controller() + from calibre.utils.terminal import colored usage = textwrap.dedent(usage) if epilog is None: - epilog = _('Created by ')+tc.RED+__author__+tc.NORMAL + epilog = _('Created by ')+colored(__author__, fg='cyan') usage += '\n\n'+_('''Whenever you pass arguments to %prog that have spaces in them, ''' '''enclose the arguments in quotation marks.''') _OptionParser.__init__(self, usage=usage, version=version, epilog=epilog, @@ -95,6 +95,21 @@ class OptionParser(_OptionParser): _("show this help message and exit") _("show program's version number and exit") + def print_usage(self, file=None): + from calibre.utils.terminal import ANSIStream + s = ANSIStream(file) + _OptionParser.print_usage(self, file=s) + + def print_help(self, file=None): + from calibre.utils.terminal import ANSIStream + s = ANSIStream(file) + _OptionParser.print_help(self, file=s) + + def print_version(self, file=None): + from calibre.utils.terminal import ANSIStream + s = ANSIStream(file) + _OptionParser.print_version(self, file=s) + def error(self, msg): if self.gui_mode: raise Exception(msg) diff --git a/src/calibre/utils/logging.py b/src/calibre/utils/logging.py index ad4a40c57e..e9132e079a 100644 --- a/src/calibre/utils/logging.py +++ b/src/calibre/utils/logging.py @@ -33,21 +33,18 @@ class ANSIStream(Stream): def __init__(self, stream=sys.stdout): Stream.__init__(self, stream) - from calibre.utils.terminfo import TerminalController - tc = TerminalController(stream) self.color = { - DEBUG: bytes(tc.GREEN), - INFO: bytes(''), - WARN: bytes(tc.YELLOW), - ERROR: bytes(tc.RED) + DEBUG: u'green', + INFO: None, + WARN: u'yellow', + ERROR: u'red', } - self.normal = bytes(tc.NORMAL) def prints(self, level, *args, **kwargs): - self.stream.write(self.color[level]) kwargs['file'] = self.stream - self._prints(*args, **kwargs) - self.stream.write(self.normal) + from calibre.utils.terminal import ColoredStream + with ColoredStream(self.stream, self.color[level]): + self._prints(*args, **kwargs) def flush(self): self.stream.flush() diff --git a/src/calibre/utils/terminal.py b/src/calibre/utils/terminal.py new file mode 100644 index 0000000000..1d8a10801a --- /dev/null +++ b/src/calibre/utils/terminal.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python +# vim:fileencoding=UTF-8:ts=4:sw=4:sta:et:sts=4:fdm=marker:ai +from __future__ import (unicode_literals, division, absolute_import, + print_function) + +__license__ = 'GPL v3' +__copyright__ = '2012, Kovid Goyal ' +__docformat__ = 'restructuredtext en' + +import os, sys, re +from itertools import izip + +from calibre.constants import iswindows + +def fmt(code): + return ('\033[%dm'%code).encode('ascii') + +RATTRIBUTES = dict( + izip(xrange(1, 9), ( + 'bold', + 'dark', + '', + 'underline', + 'blink', + '', + 'reverse', + 'concealed' + ) + )) +ATTRIBUTES = {v:fmt(k) for k, v in RATTRIBUTES.iteritems()} +del ATTRIBUTES[''] + +RBACKGROUNDS = dict( + izip(xrange(41, 48), ( + 'red', + 'green', + 'yellow', + 'blue', + 'magenta', + 'cyan', + 'white' + ), + )) +BACKGROUNDS = {v:fmt(k) for k, v in RBACKGROUNDS.iteritems()} + +RCOLORS = dict( + izip(xrange(31, 38), ( + 'red', + 'green', + 'yellow', + 'blue', + 'magenta', + 'cyan', + 'white', + ), + )) +COLORS = {v:fmt(k) for k, v in RCOLORS.iteritems()} + +RESET = fmt(0) + +if iswindows: + # From wincon.h + WCOLORS = {c:i for i, c in enumerate(( + 'black', 'blue', 'green', 'cyan', 'red', 'magenta', 'yellow', 'white'))} + + def to_flag(fg, bg, bold): + val = 0 + if bold: + val |= 0x08 + if fg in WCOLORS: + val |= WCOLORS[fg] + if bg in WCOLORS: + val |= (WCOLORS[bg] << 4) + return val + +def colored(text, fg=None, bg=None, bold=False): + prefix = [] + if fg is not None: + prefix.append(COLORS[fg]) + if bg is not None: + prefix.append(BACKGROUNDS[bg]) + if bold: + prefix.append(ATTRIBUTES['bold']) + prefix = b''.join(prefix) + suffix = RESET + if isinstance(text, type(u'')): + prefix = prefix.decode('ascii') + suffix = suffix.decode('ascii') + return prefix + text + suffix + +class Detect(object): + + def __init__(self, stream): + self.stream = stream or sys.stdout + self.isatty = getattr(self.stream, 'isatty', lambda : False)() + force_ansi = os.environ.has_key('CALIBRE_FORCE_ANSI') + if not self.isatty and force_ansi: + self.isatty = True + self.isansi = force_ansi or not iswindows + self.set_console = None + if not self.isansi: + try: + import msvcrt + self.msvcrt = msvcrt + self.file_handle = msvcrt.get_osfhandle(self.stream.fileno()) + from ctypes import windll + self.set_console = windll.kernel32.SetConsoleTextAttribute + except: + pass + +class ColoredStream(Detect): + + def __init__(self, stream=None, fg=None, bg=None, bold=False): + super(ColoredStream, self).__init__(stream) + self.fg, self.bg, self.bold = fg, bg, bold + if self.set_console is not None: + self.wval = to_flag(self.fg, self.bg, bold) + + def __enter__(self): + if not self.isatty: + return + if self.isansi: + if self.bold: + self.stream.write(ATTRIBUTES['bold']) + if self.bg is not None: + self.stream.write(BACKGROUNDS[self.bg]) + if self.fg is not None: + self.stream.write(COLORS[self.fg]) + elif self.set_console is not None: + if self.wval != 0: + self.set_console(self.file_handle, self.wval) + + def __exit__(self, *args, **kwargs): + if not self.isatty: + return + if not self.fg and not self.bg and not self.bold: + return + if self.isansi: + self.stream.write(RESET) + elif self.set_console is not None: + self.set_console(self.file_handle, WCOLORS['white']) + +class ANSIStream(Detect): + + ANSI_RE = re.compile(br'\033\[((?:\d|;)*)([a-zA-Z])') + + def __init__(self, stream=None): + super(ANSIStream, self).__init__(stream) + self.encoding = getattr(self.stream, 'encoding', 'utf-8') + + def write(self, text): + if isinstance(text, type(u'')): + text = text.encode(self.encoding, 'replace') + + if not self.isatty: + return self.strip_and_write(text) + + if self.isansi: + return self.stream.write(text) + + if not self.isansi and self.set_console is None: + return self.strip_and_write(text) + + self.write_and_convert(text) + + def strip_and_write(self, text): + self.stream.write(self.ANSI_RE.sub(b'', text)) + + def write_and_convert(self, text): + ''' + Write the given text to our wrapped stream, stripping any ANSI + sequences from the text, and optionally converting them into win32 + calls. + ''' + self.last_state = (None, None, False) + cursor = 0 + for match in self.ANSI_RE.finditer(text): + start, end = match.span() + self.write_plain_text(text, cursor, start) + self.convert_ansi(*match.groups()) + cursor = end + self.write_plain_text(text, cursor, len(text)) + + def write_plain_text(self, text, start, end): + if start < end: + self.stream.write(text[start:end]) + self.stream.flush() + + def convert_ansi(self, paramstring, command): + params = self.extract_params(paramstring) + self.call_win32(command, params) + + def extract_params(self, paramstring): + def split(paramstring): + for p in paramstring.split(b';'): + if p: + yield int(p) + return tuple(split(paramstring)) + + def call_win32(self, command, params): + if command != b'm': return + fg, bg, bold = self.last_state + + for param in params: + if param in RCOLORS: + fg = RCOLORS[param] + elif param in RBACKGROUNDS: + bg = RBACKGROUNDS[param] + elif param == 1: + bold = True + elif param == 0: + fg = 'white' + bg, bold = None, False + + self.last_state = (fg, bg, bold) + if fg or bg or bold: + self.set_console(self.file_handle, to_flag(fg, bg, bold)) + else: + self.set_console(self.file_handle, WCOLORS['white']) + +def geometry(): + try: + import curses + curses.setupterm() + except: + return 80, 80 + else: + width = curses.tigetnum('cols') or 80 + height = curses.tigetnum('lines') or 80 + return width, height + +def test(): + s = ANSIStream() + + text = [colored(t, fg=t)+'. '+colored(t, fg=t, bold=True)+'.' for t in + ('red', 'yellow', 'green', 'white', 'cyan', 'magenta', 'blue',)] + s.write('\n'.join(text)) + print() + diff --git a/src/calibre/utils/terminfo.py b/src/calibre/utils/terminfo.py deleted file mode 100644 index 074b6388f2..0000000000 --- a/src/calibre/utils/terminfo.py +++ /dev/null @@ -1,215 +0,0 @@ -__license__ = 'GPL v3' -__copyright__ = '2008, Kovid Goyal ' -import sys, re, os - -""" Get information about the terminal we are running in """ - -class TerminalController: - """ - A class that can be used to portably generate formatted output to - a terminal. - - `TerminalController` defines a set of instance variables whose - values are initialized to the control sequence necessary to - perform a given action. These can be simply included in normal - output to the terminal: - - >>> term = TerminalController() - >>> print 'This is '+term.GREEN+'green'+term.NORMAL - - Alternatively, the `render()` method can used, which replaces - '${action}' with the string required to perform 'action': - - >>> term = TerminalController() - >>> print term.render('This is ${GREEN}green${NORMAL}') - - If the terminal doesn't support a given action, then the value of - the corresponding instance variable will be set to ''. As a - result, the above code will still work on terminals that do not - support color, except that their output will not be colored. - Also, this means that you can test whether the terminal supports a - given action by simply testing the truth value of the - corresponding instance variable: - - >>> term = TerminalController() - >>> if term.CLEAR_SCREEN: - ... print 'This terminal supports clearing the screen.' - - Finally, if the width and height of the terminal are known, then - they will be stored in the `COLS` and `LINES` attributes. - """ - # Cursor movement: - BOL = '' #: Move the cursor to the beginning of the line - UP = '' #: Move the cursor up one line - DOWN = '' #: Move the cursor down one line - LEFT = '' #: Move the cursor left one char - RIGHT = '' #: Move the cursor right one char - - # Deletion: - CLEAR_SCREEN = '' #: Clear the screen and move to home position - CLEAR_EOL = '' #: Clear to the end of the line. - CLEAR_BOL = '' #: Clear to the beginning of the line. - CLEAR_EOS = '' #: Clear to the end of the screen - - # Output modes: - BOLD = '' #: Turn on bold mode - BLINK = '' #: Turn on blink mode - DIM = '' #: Turn on half-bright mode - REVERSE = '' #: Turn on reverse-video mode - NORMAL = '' #: Turn off all modes - - # Cursor display: - HIDE_CURSOR = '' #: Make the cursor invisible - SHOW_CURSOR = '' #: Make the cursor visible - - # Terminal size: - COLS = None #: Width of the terminal (None for unknown) - LINES = None #: Height of the terminal (None for unknown) - - # Foreground colors: - BLACK = BLUE = GREEN = CYAN = RED = MAGENTA = YELLOW = WHITE = '' - - # Background colors: - BG_BLACK = BG_BLUE = BG_GREEN = BG_CYAN = '' - BG_RED = BG_MAGENTA = BG_YELLOW = BG_WHITE = '' - - _STRING_CAPABILITIES = """ - BOL=cr UP=cuu1 DOWN=cud1 LEFT=cub1 RIGHT=cuf1 - CLEAR_SCREEN=clear CLEAR_EOL=el CLEAR_BOL=el1 CLEAR_EOS=ed BOLD=bold - BLINK=blink DIM=dim REVERSE=rev UNDERLINE=smul NORMAL=sgr0 - HIDE_CURSOR=cinvis SHOW_CURSOR=cnorm""".split() - _COLORS = """BLACK BLUE GREEN CYAN RED MAGENTA YELLOW WHITE""".split() - _ANSICOLORS = "BLACK RED GREEN YELLOW BLUE MAGENTA CYAN WHITE".split() - - def __init__(self, term_stream=sys.stdout): - """ - Create a `TerminalController` and initialize its attributes - with appropriate values for the current terminal. - `term_stream` is the stream that will be used for terminal - output; if this stream is not a tty, then the terminal is - assumed to be a dumb terminal (i.e., have no capabilities). - """ - # Curses isn't available on all platforms - try: import curses - except: return - - # If the stream isn't a tty, then assume it has no capabilities. - if os.environ.get('CALIBRE_WORKER', None) is not None or not hasattr(term_stream, 'isatty') or not term_stream.isatty(): return - - # Check the terminal type. If we fail, then assume that the - # terminal has no capabilities. - try: curses.setupterm() - except: return - - # Look up numeric capabilities. - self.COLS = curses.tigetnum('cols') - self.LINES = curses.tigetnum('lines') - - # Look up string capabilities. - for capability in self._STRING_CAPABILITIES: - (attrib, cap_name) = capability.split('=') - setattr(self, attrib, self._tigetstr(cap_name) or '') - - # Colors - set_fg = self._tigetstr('setf') - if set_fg: - for i,color in zip(range(len(self._COLORS)), self._COLORS): - setattr(self, color, curses.tparm(set_fg, i) or '') - set_fg_ansi = self._tigetstr('setaf') - if set_fg_ansi: - for i,color in zip(range(len(self._ANSICOLORS)), self._ANSICOLORS): - setattr(self, color, curses.tparm(set_fg_ansi, i) or '') - set_bg = self._tigetstr('setb') - if set_bg: - for i,color in zip(range(len(self._COLORS)), self._COLORS): - setattr(self, 'BG_'+color, curses.tparm(set_bg, i) or '') - set_bg_ansi = self._tigetstr('setab') - if set_bg_ansi: - for i,color in zip(range(len(self._ANSICOLORS)), self._ANSICOLORS): - setattr(self, 'BG_'+color, curses.tparm(set_bg_ansi, i) or '') - - def _tigetstr(self, cap_name): - # String capabilities can include "delays" of the form "$<2>". - # For any modern terminal, we should be able to just ignore - # these, so strip them out. - import curses - cap = curses.tigetstr(cap_name) or '' - return re.sub(r'\$<\d+>[/*]?', '', cap) - - def render(self, template): - """ - Replace each $-substitutions in the given template string with - the corresponding terminal control string (if it's defined) or - '' (if it's not). - """ - return re.sub(r'\$\$|\${\w+}', self._render_sub, template) - - def _render_sub(self, match): - s = match.group() - if s == '$$': return s - else: return getattr(self, s[2:-1]) - -####################################################################### -# Example use case: progress bar -####################################################################### - -class ProgressBar: - """ - A 3-line progress bar, which looks like:: - - Header - 20% [===========----------------------------------] - progress message - - The progress bar is colored, if the terminal supports color - output; and adjusts to the width of the terminal. - - If the terminal doesn't have the required capabilities, it uses a - simple progress bar. - """ - BAR = '%3d%% ${GREEN}[${BOLD}%s%s${NORMAL}${GREEN}]${NORMAL}\n' - HEADER = '${BOLD}${CYAN}%s${NORMAL}\n\n' - - def __init__(self, term, header, no_progress_bar = False): - self.term, self.no_progress_bar = term, no_progress_bar - self.fancy = self.term.CLEAR_EOL and self.term.UP and self.term.BOL - if self.fancy: - self.width = self.term.COLS or 75 - self.bar = term.render(self.BAR) - self.header = self.term.render(self.HEADER % header.center(self.width)) - if isinstance(self.header, unicode): - self.header = self.header.encode('utf-8') - self.cleared = 1 #: true if we haven't drawn the bar yet. - - def update(self, percent, message=''): - if isinstance(message, unicode): - message = message.encode('utf-8', 'replace') - - if self.no_progress_bar: - if message: - print message - elif self.fancy: - if self.cleared: - sys.stdout.write(self.header) - self.cleared = 0 - n = int((self.width-10)*percent) - msg = message.center(self.width) - sys.stdout.write( - self.term.BOL + self.term.UP + self.term.CLEAR_EOL + - (self.bar % (100*percent, '='*n, '-'*(self.width-10-n))) + - self.term.CLEAR_EOL + msg) - sys.stdout.flush() - else: - if not message: - print '%d%%'%(percent*100), - else: - print '%d%%'%(percent*100), message - sys.stdout.flush() - - - def clear(self): - if self.fancy and not self.cleared: - sys.stdout.write(self.term.BOL + self.term.CLEAR_EOL + - self.term.UP + self.term.CLEAR_EOL + - self.term.UP + self.term.CLEAR_EOL) - self.cleared = 1