Encapsulate all usage of colorama within ansi.py

cmd2/ansi.py

@@ -8,6 +8,9 @@
 from colorama import Fore, Back, Style
 from wcwidth import wcswidth
 
+# On Windows, filter ANSI escape codes out of text sent to stdout/stderr, and replace them with equivalent Win32 calls
+colorama.init(strip=False)
+
 # Values for allow_ansi setting
 ANSI_NEVER = 'Never'
 ANSI_TERMINAL = 'Terminal'
@@ -65,6 +68,9 @@
 BG_RESET = BG_COLORS['reset']
 RESET_ALL = Style.RESET_ALL
 
+BRIGHT = Style.BRIGHT
+NORMAL = Style.NORMAL
+
 # ANSI escape sequences not provided by colorama
 UNDERLINE_ENABLE = colorama.ansi.code_to_chars(4)
 UNDERLINE_DISABLE = colorama.ansi.code_to_chars(24)
@@ -180,3 +186,66 @@ def style(text: Any, *, fg: str = '', bg: str = '', bold: bool = False, underlin
 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')
+
+
+def async_alert_str(*, terminal_columns: int, prompt: str, line: str, cursor_offset: int, alert_msg: str) -> str:
+    """Calculate the desired string, including ANSI escape codes, for displaying an asynchronous alert message.
+
+    :param terminal_columns: terminal width (number of columns)
+    :param prompt: prompt that is displayed on the current line
+    :param line: current contents of the Readline line buffer
+    :param cursor_offset: the offset of the current cursor position within line
+    :param alert_msg: the message to display to the user
+    :return: the correct string so that the alert message appears to the user to be printed above the current line.
+    """
+    from colorama import Cursor
+    # Split the prompt lines since it can contain newline characters.
+    prompt_lines = prompt.splitlines()
+
+    # Calculate how many terminal lines are taken up by all prompt lines except for the last one.
+    # That will be included in the input lines calculations since that is where the cursor is.
+    num_prompt_terminal_lines = 0
+    for line in prompt_lines[:-1]:
+        line_width = ansi_safe_wcswidth(line)
+        num_prompt_terminal_lines += int(line_width / terminal_columns) + 1
+
+    # Now calculate how many terminal lines are take up by the input
+    last_prompt_line = prompt_lines[-1]
+    last_prompt_line_width = ansi_safe_wcswidth(last_prompt_line)
+
+    input_width = last_prompt_line_width + ansi_safe_wcswidth(line)
+
+    num_input_terminal_lines = int(input_width / terminal_columns) + 1
+
+    # Get the cursor's offset from the beginning of the first input line
+    cursor_input_offset = last_prompt_line_width + cursor_offset
+
+    # Calculate what input line the cursor is on
+    cursor_input_line = int(cursor_input_offset / terminal_columns) + 1
+
+    # Create a string that when printed will clear all input lines and display the alert
+    terminal_str = ''
+
+    # Move the cursor down to the last input line
+    if cursor_input_line != num_input_terminal_lines:
+        terminal_str += Cursor.DOWN(num_input_terminal_lines - cursor_input_line)
+
+    # Clear each line from the bottom up so that the cursor ends up on the first prompt line
+    total_lines = num_prompt_terminal_lines + num_input_terminal_lines
+    terminal_str += (colorama.ansi.clear_line() + Cursor.UP(1)) * (total_lines - 1)
+
+    # Clear the first prompt line
+    terminal_str += colorama.ansi.clear_line()
+
+    # Move the cursor to the beginning of the first prompt line and print the alert
+    terminal_str += '\r' + alert_msg
+    return terminal_str
+
+
+def set_title_str(title: str) -> str:
+    """Get the required string, including ANSI escape codes, for setting window title for the terminal.
+
+    :param title: new title for the window
+    :return string to write to sys.stderr in order to set the window title to the desired test
+    """
+    return colorama.ansi.set_title(title)

cmd2/cmd2.py

@@ -42,8 +42,6 @@
 from contextlib import redirect_stdout
 from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Type, Union
 
-import colorama
-
 from . import ansi
 from . import constants
 from . import plugin
@@ -359,9 +357,6 @@ def __init__(self, completekey: str = 'tab', stdin=None, stdout=None, *,
             except AttributeError:
                 pass
 
-        # Override whether ansi codes should be stripped from the output since cmd2 has its own logic for doing this
-        colorama.init(strip=False)
-
         # initialize plugin system
         # needs to be done before we call __init__(0)
         self._initialize_plugin_system()
@@ -3812,9 +3807,6 @@ def async_alert(self, alert_msg: str, new_prompt: Optional[str] = None) -> None:
         if not (vt100_support and self.use_rawinput):
             return
 
-        import shutil
-        from colorama import Cursor
-
         # Sanity check that can't fail if self.terminal_lock was acquired before calling this function
         if self.terminal_lock.acquire(blocking=False):
 
@@ -3838,50 +3830,10 @@ def async_alert(self, alert_msg: str, new_prompt: Optional[str] = None) -> None:
                     update_terminal = True
 
             if update_terminal:
-                # Get the size of the terminal
-                terminal_size = shutil.get_terminal_size()
-
-                # Split the prompt lines since it can contain newline characters.
-                prompt_lines = current_prompt.splitlines()
-
-                # Calculate how many terminal lines are taken up by all prompt lines except for the last one.
-                # That will be included in the input lines calculations since that is where the cursor is.
-                num_prompt_terminal_lines = 0
-                for line in prompt_lines[:-1]:
-                    line_width = ansi.ansi_safe_wcswidth(line)
-                    num_prompt_terminal_lines += int(line_width / terminal_size.columns) + 1
-
-                # Now calculate how many terminal lines are take up by the input
-                last_prompt_line = prompt_lines[-1]
-                last_prompt_line_width = ansi.ansi_safe_wcswidth(last_prompt_line)
-
-                input_width = last_prompt_line_width + ansi.ansi_safe_wcswidth(readline.get_line_buffer())
-
-                num_input_terminal_lines = int(input_width / terminal_size.columns) + 1
-
-                # Get the cursor's offset from the beginning of the first input line
-                cursor_input_offset = last_prompt_line_width + rl_get_point()
-
-                # Calculate what input line the cursor is on
-                cursor_input_line = int(cursor_input_offset / terminal_size.columns) + 1
-
-                # Create a string that when printed will clear all input lines and display the alert
-                terminal_str = ''
-
-                # Move the cursor down to the last input line
-                if cursor_input_line != num_input_terminal_lines:
-                    terminal_str += Cursor.DOWN(num_input_terminal_lines - cursor_input_line)
-
-                # Clear each line from the bottom up so that the cursor ends up on the first prompt line
-                total_lines = num_prompt_terminal_lines + num_input_terminal_lines
-                terminal_str += (colorama.ansi.clear_line() + Cursor.UP(1)) * (total_lines - 1)
-
-                # Clear the first prompt line
-                terminal_str += colorama.ansi.clear_line()
-
-                # Move the cursor to the beginning of the first prompt line and print the alert
-                terminal_str += '\r' + alert_msg
-
+                import shutil
+                terminal_str = ansi.async_alert_str(terminal_columns=shutil.get_terminal_size().columns,
+                                                    prompt=current_prompt, line=readline.get_line_buffer(),
+                                                    cursor_offset=rl_get_point(), alert_msg=alert_msg)
                 if rl_type == RlType.GNU:
                     sys.stderr.write(terminal_str)
                 elif rl_type == RlType.PYREADLINE:
@@ -3934,7 +3886,7 @@ def set_window_title(self, title: str) -> None:  # pragma: no cover
         # Sanity check that can't fail if self.terminal_lock was acquired before calling this function
         if self.terminal_lock.acquire(blocking=False):
             try:
-                sys.stderr.write(colorama.ansi.set_title(title))
+                sys.stderr.write(ansi.set_title_str(title))
             except AttributeError:
                 # Debugging in Pycharm has issues with setting terminal title
                 pass

cmd2/rl_utils.py

@@ -58,7 +58,7 @@ def enable_win_vt100(handle: HANDLE) -> bool:
 
             retVal = False
 
-            # Check if  ENABLE_VIRTUAL_TERMINAL_PROCESSING is already enabled
+            # Check if ENABLE_VIRTUAL_TERMINAL_PROCESSING is already enabled
             if (cur_mode.value & ENABLE_VIRTUAL_TERMINAL_PROCESSING) != 0:
                 retVal = True
 

examples/async_printing.py

@@ -19,7 +19,7 @@
           "Keep typing...",
           "Move that cursor...",
           "Pretty seamless, eh?",
-          "Feedback can also be given in the window title. Notice the arg count up there?",
+          "Feedback can also be given in the window title. Notice the alert count up there?",
           "You can stop and start the alerts by typing stop_alerts and start_alerts",
           "This demo will now continue to print alerts at random intervals"
           ]

tests/test_ansi.py

@@ -4,7 +4,6 @@
 Unit testing for cmd2/ansi.py module
 """
 import pytest
-from colorama import Fore, Back, Style
 
 import cmd2.ansi as ansi
 
@@ -13,14 +12,14 @@
 
 def test_strip_ansi():
     base_str = HELLO_WORLD
-    ansi_str = Fore.GREEN + base_str + Fore.RESET
+    ansi_str = ansi.style(base_str, fg='green')
     assert base_str != ansi_str
     assert base_str == ansi.strip_ansi(ansi_str)
 
 
 def test_ansi_safe_wcswidth():
     base_str = HELLO_WORLD
-    ansi_str = Fore.GREEN + base_str + Fore.RESET
+    ansi_str = ansi.style(base_str, fg='green')
     assert ansi.ansi_safe_wcswidth(ansi_str) != len(ansi_str)
 
 
@@ -32,19 +31,21 @@ def test_style_none():
 
 def test_style_fg():
     base_str = HELLO_WORLD
-    ansi_str = Fore.BLUE + base_str + Fore.RESET
-    assert ansi.style(base_str, fg='blue') == ansi_str
+    fg_color = 'blue'
+    ansi_str = ansi.FG_COLORS[fg_color] + base_str + ansi.FG_RESET
+    assert ansi.style(base_str, fg=fg_color) == ansi_str
 
 
 def test_style_bg():
     base_str = HELLO_WORLD
-    ansi_str = Back.GREEN + base_str + Back.RESET
-    assert ansi.style(base_str, bg='green') == ansi_str
+    bg_color = 'green'
+    ansi_str = ansi.BG_COLORS[bg_color] + base_str + ansi.BG_RESET
+    assert ansi.style(base_str, bg=bg_color) == ansi_str
 
 
 def test_style_bold():
     base_str = HELLO_WORLD
-    ansi_str = Style.BRIGHT + base_str + Style.NORMAL
+    ansi_str = ansi.BRIGHT + base_str + ansi.NORMAL
     assert ansi.style(base_str, bold=True) == ansi_str
 
 
@@ -56,9 +57,11 @@ def test_style_underline():
 
 def test_style_multi():
     base_str = HELLO_WORLD
-    ansi_str = Fore.BLUE + Back.GREEN + Style.BRIGHT + ansi.UNDERLINE_ENABLE + \
-               base_str + Fore.RESET + Back.RESET + Style.NORMAL + ansi.UNDERLINE_DISABLE
-    assert ansi.style(base_str, fg='blue', bg='green', bold=True, underline=True) == ansi_str
+    fg_color = 'blue'
+    bg_color = 'green'
+    ansi_str = ansi.FG_COLORS[fg_color] + ansi.BG_COLORS[bg_color] + ansi.BRIGHT + ansi.UNDERLINE_ENABLE + \
+               base_str + ansi.FG_RESET + ansi.BG_RESET + ansi.NORMAL + ansi.UNDERLINE_DISABLE
+    assert ansi.style(base_str, fg=fg_color, bg=bg_color, bold=True, underline=True) == ansi_str
 
 
 def test_style_color_not_exist():
@@ -72,7 +75,8 @@ def test_style_color_not_exist():
 
 
 def test_fg_lookup_exist():
-    assert ansi.fg_lookup('green') == Fore.GREEN
+    fg_color = 'green'
+    assert ansi.fg_lookup(fg_color) == ansi.FG_COLORS[fg_color]
 
 
 def test_fg_lookup_nonexist():
@@ -81,9 +85,28 @@ def test_fg_lookup_nonexist():
 
 
 def test_bg_lookup_exist():
-    assert ansi.bg_lookup('green') == Back.GREEN
+    bg_color = 'green'
+    assert ansi.bg_lookup(bg_color) == ansi.BG_COLORS[bg_color]
 
 
 def test_bg_lookup_nonexist():
     with pytest.raises(ValueError):
         ansi.bg_lookup('bar')
+
+
+def test_set_title_str():
+    OSC = '\033]'
+    BEL = '\007'
+    title = HELLO_WORLD
+    assert ansi.set_title_str(title) == OSC + '2;' + title + BEL
+
+
+@pytest.mark.parametrize('cols, prompt, line, cursor, msg, expected', [
+    (127, '(Cmd) ', 'help his', 12, ansi.style('Hello World!', fg='magenta'), '\x1b[2K\r\x1b[35mHello World!\x1b[39m'),
+    (127, '\n(Cmd) ', 'help ', 5, 'foo', '\x1b[2K\x1b[1A\x1b[2K\rfoo'),
+    (10, '(Cmd) ', 'help history of the american republic', 4, 'boo', '\x1b[3B\x1b[2K\x1b[1A\x1b[2K\x1b[1A\x1b[2K\x1b[1A\x1b[2K\x1b[1A\x1b[2K\rboo')
+])
+def test_async_alert_str(cols, prompt, line, cursor, msg, expected):
+    alert_str = ansi.async_alert_str(terminal_columns=cols, prompt=prompt, line=line, cursor_offset=cursor,
+                                     alert_msg=msg)
+    assert alert_str == expected