From c706f6a95c41392fcca0b3a93b689a19ba06a0f4 Mon Sep 17 00:00:00 2001
From: Kevin Van Brunt <kmvanbrunt@gmail.com>
Date: Fri, 21 Sep 2018 17:27:58 -0400
Subject: Made sure all prompts sent to GNU readline are made safe

---
 cmd2/rl_utils.py | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

(limited to 'cmd2/rl_utils.py')

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
-- 
cgit v1.2.1