From 8e996f80a452a7e5b1dfd85187dec8f5e7451aeb 2015-08-16 17:35:31
From: Aaron Meurer <asmeurer@gmail.com>
Date: 2015-08-16 17:35:31
Subject: [PATCH] Be a little smarter about invisible characters in terminal prompts

This is a partial fix to #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
index 3dd9d8c..837afc6 100644
--- 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 @@ class PromptManager(Configurable):
             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
index 7f12804..e1bcae6 100644
--- a/IPython/core/tests/test_prompts.py
+++ b/IPython/core/tests/test_prompts.py
@@ -6,7 +6,7 @@ import unittest
 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 @@ class PromptTests(unittest.TestCase):
             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)