Added optional color argument to poutput() for providing a color escape code to colorize the output with

tleonhardt · tleonhardt · commit a2ad70bd79a7 · 2018-09-24T11:13:42.000-04:00
Also:
- Added optional err_color and war_color arguments to perror() for providing color escape codes
- Removed usage of deprecated colorize() method within the examples
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
@@ -561,22 +561,25 @@ def decolorized_write(self, fileobj: IO, msg: str) -> None:
             msg = utils.strip_ansi(msg)
         fileobj.write(msg)
 
-    def poutput(self, msg: Any, end: str='\n') -> None:
+    def poutput(self, msg: Any, end: str='\n', color: str='') -> None:
         """Smarter self.stdout.write(); color aware and adds newline of not present.
 
         Also handles BrokenPipeError exceptions for when a commands's output has
         been piped to another process and that process terminates before the
         cmd2 command is finished executing.
 
         :param msg: message to print to current stdout (anything convertible to a str with '{}'.format() is OK)
-        :param end: string appended after the end of the message if not already present, default a newline
+        :param end: (optional) string appended after the end of the message if not already present, default a newline
+        :param color: (optional) color escape to output this message with
         """
         if msg is not None and msg != '':
             try:
                 msg_str = '{}'.format(msg)
-                self.decolorized_write(self.stdout, msg_str)
                 if not msg_str.endswith(end):
-                    self.decolorized_write(self.stdout, end)
+                    msg_str += end
+                if color:
+                    msg_str = color + msg_str + Fore.RESET
+                self.decolorized_write(self.stdout, msg_str)
             except BrokenPipeError:
                 # This occurs if a command's output is being piped to another
                 # process and that process closes before the command is
@@ -586,12 +589,14 @@ def poutput(self, msg: Any, end: str='\n') -> None:
                 if self.broken_pipe_warning:
                     sys.stderr.write(self.broken_pipe_warning)
 
-    def perror(self, err: Union[str, Exception], traceback_war: bool=True) -> None:
+    def perror(self, err: Union[str, Exception], traceback_war: bool=True, err_color: str=Fore.LIGHTRED_EX,
+               war_color: str=Fore.LIGHTYELLOW_EX) -> None:
         """ Print error message to sys.stderr and if debug is true, print an exception Traceback if one exists.
 
         :param err: an Exception or error message to print out
         :param traceback_war: (optional) if True, print a message to let user know they can enable debug
-        :return:
+        :param err_color: (optional) color escape to output error with
+        :param war_color: (optional) color escape to output warning with
         """
         if self.debug:
             import traceback
@@ -601,12 +606,12 @@ def perror(self, err: Union[str, Exception], traceback_war: bool=True) -> None:
             err_msg = "EXCEPTION of type '{}' occurred with message: '{}'\n".format(type(err).__name__, err)
         else:
             err_msg = "ERROR: {}\n".format(err)
-        err_msg = Fore.RED + err_msg + Fore.RESET
+        err_msg = err_color + err_msg + Fore.RESET
         self.decolorized_write(sys.stderr, err_msg)
 
         if traceback_war:
             war = "To enable full traceback, run the following command:  'set debug true'\n"
-            war = Fore.YELLOW + war + Fore.RESET
+            war = war_color + war + Fore.RESET
             self.decolorized_write(sys.stderr, war)
 
     def pfeedback(self, msg: str) -> None:
@@ -679,7 +684,7 @@ def ppaged(self, msg: str, end: str='\n', chop: bool=False) -> None:
             except BrokenPipeError:
                 # This occurs if a command's output is being piped to another process and that process closes before the
                 # command is finished. If you would like your application to print a warning message, then set the
-                # broken_pipe_warning attribute to the message you want printed.
+                # broken_pipe_warning attribute to the message you want printed.`
                 if self.broken_pipe_warning:
                     sys.stderr.write(self.broken_pipe_warning)
 
diff --git a/examples/pirate.py b/examples/pirate.py
@@ -8,19 +8,32 @@
 """
 import argparse
 
+from colorama import Fore
+
 import cmd2
 
+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,
+}
+
 
 class Pirate(cmd2.Cmd):
     """A piratical example cmd2 application involving looting and drinking."""
     def __init__(self):
         self.default_to_shell = True
         self.multiline_commands = ['sing']
         self.terminators = self.terminators + ['...']
-        self.songcolor = 'blue'
+        self.songcolor = Fore.BLUE
 
         # Add stuff to settable and/or shortcuts before calling base class initializer
-        self.settable['songcolor'] = 'Color to ``sing`` in (red/blue/green/cyan/magenta, bold, underline)'
+        self.settable['songcolor'] = 'Color to ``sing`` in (black/red/green/yellow/blue/magenta/cyan/white)'
         self.shortcuts.update({'~': 'sing'})
 
         """Initialize the base class as well as this one"""
@@ -68,7 +81,8 @@ def do_quit(self, arg):
 
     def do_sing(self, arg):
         """Sing a colorful song."""
-        self.poutput(self.colorize(arg, self.songcolor))
+        color_escape = COLORS.get(self.songcolor, default=Fore.RESET)
+        self.poutput(arg, color=color_escape)
 
     yo_parser = argparse.ArgumentParser()
     yo_parser.add_argument('--ho', type=int, default=2, help="How often to chant 'ho'")
diff --git a/examples/python_scripting.py b/examples/python_scripting.py
@@ -17,6 +17,8 @@
 import argparse
 import os
 
+from colorama import Fore
+
 import cmd2
 
 
@@ -33,7 +35,7 @@ def __init__(self):
     def _set_prompt(self):
         """Set prompt so it displays the current working directory."""
         self.cwd = os.getcwd()
-        self.prompt = self.colorize('{!r} $ '.format(self.cwd), 'cyan')
+        self.prompt = Fore.CYAN + '{!r} $ '.format(self.cwd) + Fore.RESET
 
     def postcmd(self, stop: bool, line: str) -> bool:
         """Hook method executed just after a command dispatch is finished.
diff --git a/tests/test_cmd2.py b/tests/test_cmd2.py
@@ -1094,18 +1094,19 @@ def test_ansi_prompt_not_esacped(base_app):
 
 
 def test_ansi_prompt_escaped():
+    from colorama import Fore
     app = cmd2.Cmd()
-    color = 'cyan'
+    color = Fore.CYAN
     prompt = 'InColor'
-    color_prompt = app.colorize(prompt, color)
+    color_prompt = color + prompt + Fore.RESET
 
     readline_hack_start = "\x01"
     readline_hack_end = "\x02"
 
     readline_safe_prompt = app._surround_ansi_escapes(color_prompt)
     assert prompt != color_prompt
-    assert readline_safe_prompt.startswith(readline_hack_start + app._colorcodes[color][True] + readline_hack_end)
-    assert readline_safe_prompt.endswith(readline_hack_start + app._colorcodes[color][False] + readline_hack_end)
+    assert readline_safe_prompt.startswith(readline_hack_start + color + readline_hack_end)
+    assert readline_safe_prompt.endswith(readline_hack_start + Fore.RESET + readline_hack_end)
 
 
 class HelpApp(cmd2.Cmd):
