Merge pull request #704 from python-cmd2/feature/revisit-color-support

#698 - Revisiting Color Support

Author: tleonhardt

CHANGELOG.md

@@ -2,12 +2,33 @@
 * Enhancements
     * Added support for and testing with Python 3.8, starting with 3.8 beta
     * Improved information displayed during transcript testing
+    * Added `ansi` module with functions and constants to support ANSI escape sequences which are used for things
+    like applying style to text
+    * Added support for applying styles (color, bold, underline) to text via `style()` function in `ansi` module
+    * Added default styles to ansi.py for printing `success`, `warning`. and `error` text. These are the styles used
+    by cmd2 and can be overridden to match the color scheme of your application.
+    * Added `ansi_aware_write()` function to `ansi` module. This function takes into account the value of `allow_ansi`
+    to determine if ANSI escape sequences should be stripped when not writing to a tty. See documentation for more
+    information on the `allow_ansi` setting.
 * Breaking Changes
     * Python 3.4 reached its [end of life](https://www.python.org/dev/peps/pep-0429/) on March 18, 2019 and is no longer supported by `cmd2`
         * If you need to use Python 3.4, you should pin your requirements to use `cmd2` 0.9.13
     * Made lots of changes to minimize the public API of the `cmd2.Cmd` class
         * Attributes and methods we do not intend to be public now all begin with an underscore
         * We make no API stability guarantees about these internal functions
+    * Split `perror` into 2 functions:
+        * `perror` - print a message to sys.stderr
+        * `pexcept` - print Exception message to sys.stderr. If debug is true, print exception traceback if one exists
+    * Signature of `poutput` and `perror` significantly changed
+        * Removed color parameters `color`, `err_color`, and `war_color` from `poutput` and `perror` 
+            * See the docstrings of these methods or the [cmd2 docs](https://cmd2.readthedocs.io/en/latest/unfreefeatures.html#poutput-pfeedback-perror-ppaged) for more info on applying styles to output messages
+        * `end` argument is now keyword-only and cannot be specified positionally
+        * `traceback_war` no longer exists as an argument since it isn't needed now that `perror` and `pexcept` exist
+    * Moved `cmd2.Cmd.colors` to ansi.py and renamed it to `allow_ansi`. This is now an application-wide setting.
+    * Renamed the following constants and moved them to ansi.py
+        * `COLORS_ALWAYS` --> `ANSI_ALWAYS`
+        * `COLORS_NEVER` --> `ANSI_NEVER`
+        * `COLORS_TERMINAL` --> `ANSI_TERMINAL`
 * **Renamed Commands Notice**
     * The following commands have been renamed. The old names will be supported until the next release. 
         * `load` --> `run_script`

README.md

@@ -43,6 +43,7 @@ Main Features
 - Built-in regression testing framework for your applications (transcript-based testing)
 - Transcripts for use with built-in regression can be automatically generated from `history -t` or `run_script -t`
 - Alerts that seamlessly print while user enters text at prompt
+- Colored and stylized output using `ansi.style()`
 
 Python 2.7 support is EOL
 -------------------------
@@ -89,7 +90,7 @@ Instructions for implementing each feature follow.
     class MyApp(cmd2.Cmd):
         def do_foo(self, args):
             """This docstring is the built-in help for the foo command."""
-            print('foo bar baz')
+            print(cmd2.ansi.style('foo bar baz', fg='red'))
     ```
     - By default the docstring for your **do_foo** method is the help for the **foo** command
         - NOTE: This doesn't apply if you use one of the `argparse` decorators mentioned below
@@ -314,11 +315,10 @@ example/transcript_regex.txt:
 
 ```text
 # Run this transcript with "python example.py -t transcript_regex.txt"
-# The regex for colors is because no color on Windows.
 # The regex for editor will match whatever program you use.
 # regexes on prompts just make the trailing space obvious
 (Cmd) set
-colors: /(True|False)/
+allow_ansi: Terminal
 continuation_prompt: >/ /
 debug: False
 echo: False
@@ -331,9 +331,7 @@ quiet: False
 timing: False
 ```
 
-Note how a regular expression `/(True|False)/` is used for output of the **show color** command since
-colored text is currently not available for cmd2 on Windows.  Regular expressions can be used anywhere within a
-transcript file simply by enclosing them within forward slashes, `/`.
+Regular expressions can be used anywhere within a transcript file simply by enclosing them within forward slashes, `/`.
 
 
 Found a bug?
@@ -357,12 +355,16 @@ Here are a few examples of open-source projects which use `cmd2`:
     * [Ceph](https://ceph.com/) is a distributed object, block, and file storage platform
 * [JSShell](https://github.com/Den1al/JSShell)
     * An interactive multi-user web JavaScript shell
+* [psiTurk](https://psiturk.org)
+    * An open platform for science on Amazon Mechanical Turk
 * [Jok3r](http://www.jok3r-framework.com)
     * Network & Web Pentest Automation Framework
 * [Poseidon](https://github.com/CyberReboot/poseidon)
     * Leverages software-defined networks (SDNs) to acquire and then feed network traffic to a number of machine learning techniques
 * [Unipacker](https://github.com/unipacker/unipacker)
     * Automatic and platform-independent unpacker for Windows binaries based on emulation
+* [FLASHMINGO](https://github.com/fireeye/flashmingo)
+    * Automatic analysis of SWF files based on some heuristics. Extensible via plugins.
 * [tomcatmanager](https://github.com/tomcatmanager/tomcatmanager)
     * A command line tool and python library for managing a tomcat server
 * [mptcpanalyzer](https://github.com/teto/mptcpanalyzer)

cmd2/ansi.py

@@ -0,0 +1,182 @@
+# coding=utf-8
+"""Support for ANSI escape sequences which are used for things like applying style to text"""
+import functools
+import re
+from typing import Any, IO
+
+import colorama
+from colorama import Fore, Back, Style
+from wcwidth import wcswidth
+
+# Values for allow_ansi setting
+ANSI_NEVER = 'Never'
+ANSI_TERMINAL = 'Terminal'
+ANSI_ALWAYS = 'Always'
+
+# Controls when ANSI escape sequences are allowed in output
+allow_ansi = ANSI_TERMINAL
+
+# Regular expression to match ANSI escape sequences
+ANSI_ESCAPE_RE = re.compile(r'\x1b[^m]*m')
+
+# Foreground color presets
+FG_COLORS = {
+    'black': Fore.BLACK,
+    'red': Fore.RED,
+    'green': Fore.GREEN,
+    'yellow': Fore.YELLOW,
+    'blue': Fore.BLUE,
+    'magenta': Fore.MAGENTA,
+    'cyan': Fore.CYAN,
+    'white': Fore.WHITE,
+    'bright_black': Fore.LIGHTBLACK_EX,
+    'bright_red': Fore.LIGHTRED_EX,
+    'bright_green': Fore.LIGHTGREEN_EX,
+    'bright_yellow': Fore.LIGHTYELLOW_EX,
+    'bright_blue': Fore.LIGHTBLUE_EX,
+    'bright_magenta': Fore.LIGHTMAGENTA_EX,
+    'bright_cyan': Fore.LIGHTCYAN_EX,
+    'bright_white': Fore.LIGHTWHITE_EX,
+    'reset': Fore.RESET,
+}
+
+# Background color presets
+BG_COLORS = {
+    'black': Back.BLACK,
+    'red': Back.RED,
+    'green': Back.GREEN,
+    'yellow': Back.YELLOW,
+    'blue': Back.BLUE,
+    'magenta': Back.MAGENTA,
+    'cyan': Back.CYAN,
+    'white': Back.WHITE,
+    'bright_black': Back.LIGHTBLACK_EX,
+    'bright_red': Back.LIGHTRED_EX,
+    'bright_green': Back.LIGHTGREEN_EX,
+    'bright_yellow': Back.LIGHTYELLOW_EX,
+    'bright_blue': Back.LIGHTBLUE_EX,
+    'bright_magenta': Back.LIGHTMAGENTA_EX,
+    'bright_cyan': Back.LIGHTCYAN_EX,
+    'bright_white': Back.LIGHTWHITE_EX,
+    'reset': Back.RESET,
+}
+
+FG_RESET = FG_COLORS['reset']
+BG_RESET = BG_COLORS['reset']
+RESET_ALL = Style.RESET_ALL
+
+# ANSI escape sequences not provided by colorama
+UNDERLINE_ENABLE = colorama.ansi.code_to_chars(4)
+UNDERLINE_DISABLE = colorama.ansi.code_to_chars(24)
+
+
+def strip_ansi(text: str) -> str:
+    """
+    Strip ANSI escape sequences from a string.
+
+    :param text: string which may contain ANSI escape sequences
+    :return: the same string with any ANSI escape sequences removed
+    """
+    return ANSI_ESCAPE_RE.sub('', text)
+
+
+def ansi_safe_wcswidth(text: str) -> int:
+    """
+    Wrap wcswidth to make it compatible with strings that contains ANSI escape sequences
+
+    :param text: the string being measured
+    """
+    # Strip ANSI escape sequences since they cause wcswidth to return -1
+    return wcswidth(strip_ansi(text))
+
+
+def ansi_aware_write(fileobj: IO, msg: str) -> None:
+    """
+    Write a string to a fileobject and strip its ANSI escape sequences if required by allow_ansi setting
+
+    :param fileobj: the file object being written to
+    :param msg: the string being written
+    """
+    if allow_ansi.lower() == ANSI_NEVER.lower() or \
+            (allow_ansi.lower() == ANSI_TERMINAL.lower() and not fileobj.isatty()):
+        msg = strip_ansi(msg)
+    fileobj.write(msg)
+
+
+def fg_lookup(fg_name: str) -> str:
+    """Look up ANSI escape codes based on foreground color name.
+
+    :param fg_name: foreground color name to look up ANSI escape code(s) for
+    :return: ANSI escape code(s) associated with this color
+    :raises ValueError if the color cannot be found
+    """
+    try:
+        ansi_escape = FG_COLORS[fg_name.lower()]
+    except KeyError:
+        raise ValueError('Foreground color {!r} does not exist.'.format(fg_name))
+    return ansi_escape
+
+
+def bg_lookup(bg_name: str) -> str:
+    """Look up ANSI escape codes based on background color name.
+
+    :param bg_name: background color name to look up ANSI escape code(s) for
+    :return: ANSI escape code(s) associated with this color
+    :raises ValueError if the color cannot be found
+    """
+    try:
+        ansi_escape = BG_COLORS[bg_name.lower()]
+    except KeyError:
+        raise ValueError('Background color {!r} does not exist.'.format(bg_name))
+    return ansi_escape
+
+
+def style(text: Any, *, fg: str = '', bg: str = '', bold: bool = False, underline: bool = False) -> str:
+    """Styles a string with ANSI colors and/or styles and returns the new string.
+
+    The styling is self contained which means that at the end of the string reset code(s) are issued
+    to undo whatever styling was done at the beginning.
+
+    :param text: Any object compatible with str.format()
+    :param fg: foreground color. Relies on `fg_lookup()` to retrieve ANSI escape based on name. Defaults to no color.
+    :param bg: background color. Relies on `bg_lookup()` to retrieve ANSI escape based on name. Defaults to no color.
+    :param bold: apply the bold style if True. Defaults to False.
+    :param underline: apply the underline style if True. Defaults to False.
+    :return: the stylized string
+    """
+    # List of strings that add style
+    additions = []
+
+    # List of strings that remove style
+    removals = []
+
+    # Convert the text object into a string if it isn't already one
+    text = "{}".format(text)
+
+    # Process the style settings
+    if fg:
+        additions.append(fg_lookup(fg))
+        removals.append(FG_RESET)
+
+    if bg:
+        additions.append(bg_lookup(bg))
+        removals.append(BG_RESET)
+
+    if bold:
+        additions.append(Style.BRIGHT)
+        removals.append(Style.NORMAL)
+
+    if underline:
+        additions.append(UNDERLINE_ENABLE)
+        removals.append(UNDERLINE_DISABLE)
+
+    # Combine the ANSI escape sequences with the text
+    return "".join(additions) + text + "".join(removals)
+
+
+# Default styles for printing strings of various types.
+# These can be altered to suit an application's needs and only need to be a
+# function with the following structure: func(str) -> str
+style_success = functools.partial(style, fg='green', bold=True)
+style_warning = functools.partial(style, fg='bright_yellow')
+style_error = functools.partial(style, fg='bright_red')

cmd2/argparse_completer.py

@@ -66,10 +66,8 @@ def my_completer(text: str, line: str, begidx: int, endidx:int, extra_param: str
 from argparse import ZERO_OR_MORE, ONE_OR_MORE, ArgumentError, _, _get_action_name, SUPPRESS
 from typing import List, Dict, Tuple, Callable, Union
 
-from colorama import Fore
-
+from .ansi import ansi_aware_write, ansi_safe_wcswidth, style_error
 from .rl_utils import rl_force_redisplay
-from .utils import ansi_safe_wcswidth
 
 # attribute that can optionally added to an argparse argument (called an Action) to
 # define the completion choices for the argument. You may provide a Collection or a Function.
@@ -996,7 +994,8 @@ def error(self, message: str) -> None:
             linum += 1
 
         self.print_usage(sys.stderr)
-        self.exit(2, Fore.LIGHTRED_EX + '{}\n\n'.format(formatted_message) + Fore.RESET)
+        formatted_message = style_error(formatted_message)
+        self.exit(2, '{}\n\n'.format(formatted_message))
 
     def format_help(self) -> str:
         """Copy of format_help() from argparse.ArgumentParser with tweaks to separately display required parameters"""
@@ -1048,6 +1047,13 @@ def format_help(self) -> str:
         # determine help from format above
         return formatter.format_help() + '\n'
 
+    def _print_message(self, message, file=None):
+        # Override _print_message to use ansi_aware_write() since we use ANSI escape characters to support color
+        if message:
+            if file is None:
+                file = _sys.stderr
+            ansi_aware_write(file, message)
+
     def _get_nargs_pattern(self, action) -> str:
         # Override _get_nargs_pattern behavior to use the nargs ranges provided by AutoCompleter
         if isinstance(action, _RangeAction) and \