Updated async_alert() to account for self.prompt not matching Readline's current prompt.async_prompt

4 files changed, 87 insertions, 51 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3c8f9403..f65a3767 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,7 @@
   * Removed `--verbose` flag from set command since descriptions always show now.
   * All cmd2 built-in commands now populate `self.last_result`.
   * Argparse tab completer will complete remaining flag names if there are no more positionals to complete.
+  * Updated `async_alert()` to account for `self.prompt` not matching Readline's current prompt.
 * Deletions (potentially breaking changes)
   * Deleted ``set_choices_provider()`` and ``set_completer()`` which were deprecated in 2.1.2
 
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index eaa0655d..b787bb18 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -124,8 +124,9 @@ from .parsing import (
 )
 from .rl_utils import (
     RlType,
+    rl_escape_prompt,
     rl_get_point,
-    rl_make_safe_prompt,
+    rl_get_prompt,
     rl_set_prompt,
     rl_type,
     rl_warning,
@@ -2982,11 +2983,11 @@ class Cmd(cmd.Cmd):
             if sys.stdin.isatty():
                 try:
                     # Deal with the vagaries of readline and ANSI escape codes
-                    safe_prompt = rl_make_safe_prompt(prompt)
+                    escaped_prompt = rl_escape_prompt(prompt)
 
                     with self.sigint_protection:
                         configure_readline()
-                    line = input(safe_prompt)
+                    line = input(escaped_prompt)
                 finally:
                     with self.sigint_protection:
                         restore_readline()
@@ -5013,12 +5014,12 @@ class Cmd(cmd.Cmd):
         Raises a `RuntimeError` if called while another thread holds `terminal_lock`.
 
         IMPORTANT: This function will not print an alert unless it can acquire self.terminal_lock to ensure
-                   a prompt is onscreen.  Therefore it is best to acquire the lock before calling this function
+                   a prompt is onscreen. Therefore it is best to acquire the lock before calling this function
                    to guarantee the alert prints and to avoid raising a RuntimeError.
 
         :param alert_msg: the message to display to the user
-        :param new_prompt: if you also want to change the prompt that is displayed, then include it here
-                           see async_update_prompt() docstring for guidance on updating a prompt
+        :param new_prompt: If you also want to change the prompt that is displayed, then include it here.
+                           See async_update_prompt() docstring for guidance on updating a prompt.
         """
         if not (vt100_support and self.use_rawinput):
             return
@@ -5026,29 +5027,31 @@ class Cmd(cmd.Cmd):
         # Sanity check that can't fail if self.terminal_lock was acquired before calling this function
         if self.terminal_lock.acquire(blocking=False):
 
-            # Only update terminal if there are changes
+            # Windows terminals tend to flicker when we redraw the prompt and input lines.
+            # To reduce how often this occurs, only update terminal if there are changes.
             update_terminal = False
 
             if alert_msg:
                 alert_msg += '\n'
                 update_terminal = True
 
-            # Set the prompt if it's changed
-            if new_prompt is not None and new_prompt != self.prompt:
+            if new_prompt is not None:
                 self.prompt = new_prompt
 
-                # If we aren't at a continuation prompt, then it's OK to update it
-                if not self._at_continuation_prompt:
-                    rl_set_prompt(self.prompt)
-                    update_terminal = True
+            # Check if the prompt to display has changed from what's currently displayed
+            cur_onscreen_prompt = rl_get_prompt()
+            new_onscreen_prompt = self.continuation_prompt if self._at_continuation_prompt else self.prompt
+
+            if new_onscreen_prompt != cur_onscreen_prompt:
+                update_terminal = True
 
             if update_terminal:
                 import shutil
 
-                current_prompt = self.continuation_prompt if self._at_continuation_prompt else self.prompt
+                # Generate the string which will replace the current prompt and input lines with the alert
                 terminal_str = ansi.async_alert_str(
                     terminal_columns=shutil.get_terminal_size().columns,
-                    prompt=current_prompt,
+                    prompt=cur_onscreen_prompt,
                     line=readline.get_line_buffer(),
                     cursor_offset=rl_get_point(),
                     alert_msg=alert_msg,
@@ -5060,7 +5063,10 @@ class Cmd(cmd.Cmd):
                     # noinspection PyUnresolvedReferences
                     readline.rl.mode.console.write(terminal_str)
 
-                # Redraw the prompt and input lines
+                # Update Readline's prompt before we redraw it
+                rl_set_prompt(new_onscreen_prompt)
+
+                # Redraw the prompt and input lines below the alert
                 rl_force_redisplay()
 
             self.terminal_lock.release()
@@ -5079,7 +5085,7 @@ class Cmd(cmd.Cmd):
         Raises a `RuntimeError` if called while another thread holds `terminal_lock`.
 
         IMPORTANT: This function will not update the prompt unless it can acquire self.terminal_lock to ensure
-                   a prompt is onscreen.  Therefore it is best to acquire the lock before calling this function
+                   a prompt is onscreen. Therefore it is best to acquire the lock before calling this function
                    to guarantee the prompt changes and to avoid raising a RuntimeError.
 
                    If user is at a continuation prompt while entering a multiline command, the onscreen prompt will
diff --git a/cmd2/rl_utils.py b/cmd2/rl_utils.py
index a79f1519..b2dc7649 100644
--- a/cmd2/rl_utils.py
+++ b/cmd2/rl_utils.py
@@ -1,11 +1,15 @@
 # coding=utf-8
 """
-Imports the proper readline for the platform and provides utility functions for it
+Imports the proper Readline for the platform and provides utility functions for it
 """
 import sys
 from enum import (
     Enum,
 )
+from typing import (
+    Union,
+    cast,
+)
 
 # Prefer statically linked gnureadline if available (for macOS compatibility due to issues with libedit)
 try:
@@ -29,13 +33,13 @@ class RlType(Enum):
     NONE = 3
 
 
-# Check what implementation of readline we are using
+# Check what implementation of Readline we are using
 rl_type = RlType.NONE
 
 # Tells if the terminal we are running in supports vt100 control characters
 vt100_support = False
 
-# Explanation for why readline wasn't loaded
+# Explanation for why Readline wasn't loaded
 _rl_warn_reason = ''
 
 # The order of this check matters since importing pyreadline/pyreadline3 will also show readline in the modules list
@@ -188,23 +192,43 @@ def rl_get_point() -> int:  # pragma: no cover
         return 0
 
 
-# noinspection PyProtectedMember, PyUnresolvedReferences
+# noinspection PyUnresolvedReferences
+def rl_get_prompt() -> str:  # pragma: no cover
+    """Gets Readline's current prompt"""
+    if rl_type == RlType.GNU:
+        encoded_prompt = ctypes.c_char_p.in_dll(readline_lib, "rl_prompt").value
+        prompt = cast(bytes, encoded_prompt).decode(encoding='utf-8')
+
+    elif rl_type == RlType.PYREADLINE:
+        prompt_data: Union[str, bytes] = readline.rl.prompt
+        if isinstance(prompt_data, bytes):
+            prompt = prompt_data.decode(encoding='utf-8')
+        else:
+            prompt = prompt_data
+
+    else:
+        prompt = ''
+
+    return rl_unescape_prompt(prompt)
+
+
+# noinspection PyUnresolvedReferences
 def rl_set_prompt(prompt: str) -> None:  # pragma: no cover
     """
-    Sets readline's prompt
+    Sets Readline's prompt
     :param prompt: the new prompt value
     """
-    safe_prompt = rl_make_safe_prompt(prompt)
+    escaped_prompt = rl_escape_prompt(prompt)
 
     if rl_type == RlType.GNU:
-        encoded_prompt = bytes(safe_prompt, encoding='utf-8')
+        encoded_prompt = bytes(escaped_prompt, encoding='utf-8')
         readline_lib.rl_set_prompt(encoded_prompt)
 
     elif rl_type == RlType.PYREADLINE:
-        readline.rl._set_prompt(safe_prompt)
+        readline.rl.prompt = escaped_prompt
 
 
-def rl_make_safe_prompt(prompt: str) -> str:  # pragma: no cover
+def rl_escape_prompt(prompt: str) -> str:
     """Overcome bug in GNU Readline in relation to calculation of prompt length in presence of ANSI escape codes
 
     :param prompt: original prompt
@@ -212,20 +236,20 @@ def rl_make_safe_prompt(prompt: str) -> str:  # pragma: no cover
     """
     if rl_type == RlType.GNU:
         # start code to tell GNU Readline about beginning of invisible characters
-        start = "\x01"
+        escape_start = "\x01"
 
         # end code to tell GNU Readline about end of invisible characters
-        end = "\x02"
+        escape_end = "\x02"
 
         escaped = False
         result = ""
 
         for c in prompt:
             if c == "\x1b" and not escaped:
-                result += start + c
+                result += escape_start + c
                 escaped = True
             elif c.isalpha() and escaped:
-                result += c + end
+                result += c + escape_end
                 escaped = False
             else:
                 result += c
@@ -234,3 +258,13 @@ def rl_make_safe_prompt(prompt: str) -> str:  # pragma: no cover
 
     else:
         return prompt
+
+
+def rl_unescape_prompt(prompt: str) -> str:
+    """Remove escape characters from a Readline prompt"""
+    if rl_type == RlType.GNU:
+        escape_start = "\x01"
+        escape_end = "\x02"
+        prompt = prompt.replace(escape_start, "").replace(escape_end, "")
+
+    return prompt
diff --git a/tests/test_cmd2.py b/tests/test_cmd2.py
index e1a52bce..0f022849 100755
--- a/tests/test_cmd2.py
+++ b/tests/test_cmd2.py
@@ -1020,37 +1020,32 @@ def test_default_to_shell(base_app, monkeypatch):
     assert m.called
 
 
-def test_ansi_prompt_not_esacped(base_app):
+def test_escaping_prompt():
     from cmd2.rl_utils import (
-        rl_make_safe_prompt,
+        rl_escape_prompt,
+        rl_unescape_prompt,
     )
 
+    # This prompt has nothing which needs to be escaped
     prompt = '(Cmd) '
-    assert rl_make_safe_prompt(prompt) == prompt
+    assert rl_escape_prompt(prompt) == prompt
 
-
-def test_ansi_prompt_escaped():
-    from cmd2.rl_utils import (
-        rl_make_safe_prompt,
-    )
-
-    app = cmd2.Cmd()
+    # This prompt has color which needs to be escaped
     color = 'cyan'
-    prompt = 'InColor'
-    color_prompt = ansi.style(prompt, fg=color)
+    prompt = ansi.style('InColor', fg=color)
 
-    readline_hack_start = "\x01"
-    readline_hack_end = "\x02"
+    escape_start = "\x01"
+    escape_end = "\x02"
 
-    readline_safe_prompt = rl_make_safe_prompt(color_prompt)
-    assert prompt != color_prompt
+    escaped_prompt = rl_escape_prompt(prompt)
     if sys.platform.startswith('win'):
-        # PyReadline on Windows doesn't suffer from the GNU readline bug which requires the hack
-        assert readline_safe_prompt.startswith(ansi.fg_lookup(color))
-        assert readline_safe_prompt.endswith(ansi.FG_RESET)
+        # PyReadline on Windows doesn't need to escape invisible characters
+        assert escaped_prompt == prompt
     else:
-        assert readline_safe_prompt.startswith(readline_hack_start + ansi.fg_lookup(color) + readline_hack_end)
-        assert readline_safe_prompt.endswith(readline_hack_start + ansi.FG_RESET + readline_hack_end)
+        assert escaped_prompt.startswith(escape_start + ansi.fg_lookup(color) + escape_end)
+        assert escaped_prompt.endswith(escape_start + ansi.FG_RESET + escape_end)
+
+    assert rl_unescape_prompt(escaped_prompt) == prompt
 
 
 class HelpApp(cmd2.Cmd):