Merge pull request ipython#8738 from asmeurer/invisible-characters

Carreau · Carreau · commit b0cddd39dadd · 2015-08-25T10:08:49.000+02:00
Be a little smarter about invisible characters in terminal prompts
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 [1]: \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)
