Add textwrap.wrap new ignore_ansi_escape parameter, to ignore ANSI escape codes

Author: alexprengere

Lib/argparse.py

@@ -691,14 +691,15 @@ def _split_lines(self, text, width):
         # The textwrap module is used only for formatting help.
         # Delay its import for speeding up the common usage of argparse.
         import textwrap
-        return textwrap.wrap(text, width)
+        return textwrap.wrap(text, width, ignore_ansi_escape=True)
 
     def _fill_text(self, text, width, indent):
         text = self._whitespace_matcher.sub(' ', text).strip()
         import textwrap
         return textwrap.fill(text, width,
                              initial_indent=indent,
-                             subsequent_indent=indent)
+                             subsequent_indent=indent,
+                             ignore_ansi_escape=True)
 
     def _get_help_string(self, action):
         return action.help

Lib/textwrap.py

@@ -61,6 +61,8 @@ class TextWrapper:
         Truncate wrapped lines.
       placeholder (default: ' [...]')
         Append to the last line of truncated text.
+      ignore_ansi_escape (default: false)
+        Ignore ANSI escape sequences when computing lengths of lines.
     """
 
     unicode_whitespace_trans = dict.fromkeys(map(ord, _whitespace), ord(' '))
@@ -109,6 +111,8 @@ class TextWrapper:
                                  r'[\"\']?'           # optional end-of-quote
                                  r'\z')               # end of chunk
 
+    ansi_escape_re = re.compile(r'\x1b\[[0-9;]*m')
+
     def __init__(self,
                  width=70,
                  initial_indent="",
@@ -122,7 +126,8 @@ def __init__(self,
                  tabsize=8,
                  *,
                  max_lines=None,
-                 placeholder=' [...]'):
+                 placeholder=' [...]',
+                 ignore_ansi_escape=False):
         self.width = width
         self.initial_indent = initial_indent
         self.subsequent_indent = subsequent_indent
@@ -135,6 +140,7 @@ def __init__(self,
         self.tabsize = tabsize
         self.max_lines = max_lines
         self.placeholder = placeholder
+        self.ignore_ansi_escape = ignore_ansi_escape
 
 
     # -- Private methods -----------------------------------------------
@@ -235,6 +241,10 @@ def _handle_long_word(self, reversed_chunks, cur_line, cur_len, width):
         # cur_len will be zero, so the next line will be entirely
         # devoted to the long word that we can't handle right now.
 
+    def _str_len_without_ansi_escape_codes(self, s):
+        """Return the length of string s without ANSI escape codes."""
+        return len(self.ansi_escape_re.sub(s, ""))
+
     def _wrap_chunks(self, chunks):
         """_wrap_chunks(chunks : [string]) -> [string]
 
@@ -259,6 +269,11 @@ def _wrap_chunks(self, chunks):
             if len(indent) + len(self.placeholder.lstrip()) > self.width:
                 raise ValueError("placeholder too large for max width")
 
+        if self.ignore_ansi_escape:
+            _str_len = self._str_len_without_ansi_escape_codes
+        else:
+            _str_len = len
+
         # Arrange in reverse order so items can be efficiently popped
         # from a stack of chucks.
         chunks.reverse()
@@ -285,7 +300,7 @@ def _wrap_chunks(self, chunks):
                 del chunks[-1]
 
             while chunks:
-                l = len(chunks[-1])
+                l = _str_len(chunks[-1])
 
                 # Can at least squeeze this chunk onto the current line.
                 if cur_len + l <= width: