Made sure all prompts sent to GNU readline are made safe

3 files changed, 42 insertions, 39 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 62077427..dae934b8 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -51,7 +51,7 @@ from .clipboard import can_clip, get_paste_buffer, write_to_paste_buffer
 from .parsing import StatementParser, Statement
 
 # Set up readline
-from .rl_utils import rl_type, RlType, rl_get_point, rl_set_prompt, vt100_support
+from .rl_utils import rl_type, RlType, rl_get_point, rl_set_prompt, vt100_support, rl_make_safe_prompt
 if rl_type == RlType.NONE:  # pragma: no cover
     rl_warning = "Readline features including tab completion have been disabled since no \n" \
                  "supported version of readline was found. To resolve this, install \n" \
@@ -2070,34 +2070,6 @@ class Cmd(cmd.Cmd):
         # Print out a message stating this is an unknown command
         self.poutput('*** Unknown syntax: {}\n'.format(arg))
 
-    @staticmethod
-    def _surround_ansi_escapes(prompt: str, start: str="\x01", end: str="\x02") -> str:
-        """Overcome bug in GNU Readline in relation to calculation of prompt length in presence of ANSI escape codes.
-
-        :param prompt: original prompt
-        :param start: start code to tell GNU Readline about beginning of invisible characters
-        :param end: end code to tell GNU Readline about end of invisible characters
-        :return: prompt safe to pass to GNU Readline
-        """
-        # Windows terminals don't use ANSI escape codes and Windows readline isn't based on GNU Readline
-        if sys.platform == "win32":
-            return prompt
-
-        escaped = False
-        result = ""
-
-        for c in prompt:
-            if c == "\x1b" and not escaped:
-                result += start + c
-                escaped = True
-            elif c.isalpha() and escaped:
-                result += c + end
-                escaped = False
-            else:
-                result += c
-
-        return result
-
     def pseudo_raw_input(self, prompt: str) -> str:
         """Began life as a copy of cmd's cmdloop; like raw_input but
 
@@ -2106,9 +2078,6 @@ class Cmd(cmd.Cmd):
           to decide whether to print the prompt and the input
         """
 
-        # Deal with the vagaries of readline and ANSI escape codes
-        safe_prompt = self._surround_ansi_escapes(prompt)
-
         if self.use_rawinput:
             try:
                 if sys.stdin.isatty():
@@ -2118,11 +2087,13 @@ class Cmd(cmd.Cmd):
                     except RuntimeError:
                         pass
 
+                    # Deal with the vagaries of readline and ANSI escape codes
+                    safe_prompt = rl_make_safe_prompt(prompt)
                     line = input(safe_prompt)
                 else:
                     line = input()
                     if self.echo:
-                        sys.stdout.write('{}{}\n'.format(safe_prompt, line))
+                        sys.stdout.write('{}{}\n'.format(self.prompt, line))
             except EOFError:
                 line = 'eof'
             finally:
@@ -2132,7 +2103,7 @@ class Cmd(cmd.Cmd):
         else:
             if self.stdin.isatty():
                 # on a tty, print the prompt first, then read the line
-                self.poutput(safe_prompt, end='')
+                self.poutput(self.prompt, end='')
                 self.stdout.flush()
                 line = self.stdin.readline()
                 if len(line) == 0:
@@ -2145,7 +2116,7 @@ class Cmd(cmd.Cmd):
                 if len(line):
                     # we read something, output the prompt and the something
                     if self.echo:
-                        self.poutput('{}{}'.format(safe_prompt, line))
+                        self.poutput('{}{}'.format(self.prompt, line))
                 else:
                     line = 'eof'
         return line.strip()
diff --git a/cmd2/rl_utils.py b/cmd2/rl_utils.py
index 96a74b67..569ba8cf 100644
--- a/cmd2/rl_utils.py
+++ b/cmd2/rl_utils.py
@@ -164,9 +164,39 @@ def rl_set_prompt(prompt: str) -> None:
     Sets readline's prompt
     :param prompt: the new prompt value
     """
+    safe_prompt = rl_make_safe_prompt(prompt)
+
     if rl_type == RlType.GNU:  # pragma: no cover
-        encoded_prompt = bytes(prompt, encoding='utf-8')
+        encoded_prompt = bytes(safe_prompt, encoding='utf-8')
         readline_lib.rl_set_prompt(encoded_prompt)
 
     elif rl_type == RlType.PYREADLINE:  # pragma: no cover
-        readline.rl._set_prompt(prompt)
+        readline.rl._set_prompt(safe_prompt)
+
+
+def rl_make_safe_prompt(prompt: str, start: str = "\x01", end: str = "\x02") -> str:
+    """Overcome bug in GNU Readline in relation to calculation of prompt length in presence of ANSI escape codes.
+
+    :param prompt: original prompt
+    :param start: start code to tell GNU Readline about beginning of invisible characters
+    :param end: end code to tell GNU Readline about end of invisible characters
+    :return: prompt safe to pass to GNU Readline
+    """
+    if rl_type == RlType.GNU:
+        escaped = False
+        result = ""
+
+        for c in prompt:
+            if c == "\x1b" and not escaped:
+                result += start + c
+                escaped = True
+            elif c.isalpha() and escaped:
+                result += c + end
+                escaped = False
+            else:
+                result += c
+
+        return result
+
+    else:
+        return prompt
diff --git a/tests/test_cmd2.py b/tests/test_cmd2.py
index 1e7e2c3f..d3d7d8f1 100644
--- a/tests/test_cmd2.py
+++ b/tests/test_cmd2.py
@@ -1093,11 +1093,13 @@ def test_default_to_shell_failure(capsys):
 
 
 def test_ansi_prompt_not_esacped(base_app):
+    from cmd2.rl_utils import rl_make_safe_prompt
     prompt = '(Cmd) '
-    assert base_app._surround_ansi_escapes(prompt) == prompt
+    assert rl_make_safe_prompt(prompt) == prompt
 
 
 def test_ansi_prompt_escaped():
+    from cmd2.rl_utils import rl_make_safe_prompt
     app = cmd2.Cmd()
     color = 'cyan'
     prompt = 'InColor'
@@ -1106,7 +1108,7 @@ def test_ansi_prompt_escaped():
     readline_hack_start = "\x01"
     readline_hack_end = "\x02"
 
-    readline_safe_prompt = app._surround_ansi_escapes(color_prompt)
+    readline_safe_prompt = rl_make_safe_prompt(color_prompt)
     if sys.platform.startswith('win'):
         # colorize() does nothing on Windows due to lack of ANSI color support
         assert prompt == color_prompt