Be a little smarter about invisible characters in terminal prompts

asmeurer · asmeurer · commit 8e996f80a452 · 2015-08-16T11:35:31.000-06:00
This is a partial fix to ipython#8724. Previously, only known color codes were considered to be invisible. Now, it looks for any kind of invisible sequence as defined by the \001 \002 delimiters (which is what readline uses). The situation could still be improved, as it still assumes that the number of invisible characters is constant for a given template. Making this work correctly with the existing API is awkward, so I didn't attempt it, especially since the readline frontend may be removed at some point in the near future.
diff --git a/IPython/core/prompts.py b/IPython/core/prompts.py
@@ -257,6 +257,14 @@ def _lenlastline(s):
     return len(s.splitlines()[-1])
 
 
+invisible_chars_re = re.compile('\001[^\001\002]*\002')
+def _invisible_characters(s):
+    """
+    Get the number of invisible ANSI characters in s. Invisible characters
+    must be delimited by \001 and \002.
+    """
+    return _lenlastline(s) - _lenlastline(invisible_chars_re.sub('', s))
+
 class UserNSFormatter(Formatter):
     """A Formatter that falls back on a shell's user_ns and __builtins__ for name resolution"""
     def __init__(self, shell):
@@ -350,8 +358,7 @@ def update_prompt(self, name, new_template=None):
             self.templates[name] = multiple_replace(prompt_abbreviations, new_template)
         # We count invisible characters (colour escapes) on the last line of the
         # prompt, to calculate the width for lining up subsequent prompts.
-        invis_chars = _lenlastline(self._render(name, color=True)) - \
-                        _lenlastline(self._render(name, color=False))
+        invis_chars = _invisible_characters(self._render(name, color=True))
         self.invisible_chars[name] = invis_chars
     
     def _update_prompt_trait(self, traitname, new_template):
diff --git a/IPython/core/tests/test_prompts.py b/IPython/core/tests/test_prompts.py
@@ -6,7 +6,7 @@
 import os
 
 from IPython.testing import tools as tt, decorators as dec
-from IPython.core.prompts import PromptManager, LazyEvaluate
+from IPython.core.prompts import PromptManager, LazyEvaluate, _invisible_characters
 from IPython.testing.globalipapp import get_ipython
 from IPython.utils.tempdir import TemporaryWorkingDirectory
 from IPython.utils import py3compat
@@ -106,4 +106,24 @@ def test_cwd_x(self):
             self.assertEqual(p, '~')
         finally:
             os.chdir(save)
-    
+
+    def test_invisible_chars(self):
+        self.assertEqual(_invisible_characters('abc'), 0)
+        self.assertEqual(_invisible_characters('\001\033[1;37m\002'), 9)
+        # Sequences must be between \001 and \002 to be counted
+        self.assertEqual(_invisible_characters('\033[1;37m'), 0)
+        # Test custom escape sequences
+        self.assertEqual(_invisible_characters('\001\033]133;A\a\002'), 10)
+
+    def test_width(self):
+        default_in = '\x01\x1b]133;A\x07\x02In [\\#]: \x01\x1b]133;B\x07\x02'
+        self.pm.in_template = default_in
+        self.pm.render('in')
+        self.assertEqual(self.pm.width, 8)
+        self.assertEqual(self.pm.txtwidth, 8)
+
+        # Test custom escape sequences
+        self.pm.in_template = '\001\033]133;A\a\002' + default_in + '\001\033]133;B\a\002'
+        self.pm.render('in')
+        self.assertEqual(self.pm.width, 8)
+        self.assertEqual(self.pm.txtwidth, 8)
