From 89f305d67e913ea1512e2618a0375359ec925ada Mon Sep 17 00:00:00 2001
From: Scott Gasch <scott@gasch.org>
Date: Wed, 23 Feb 2022 12:43:37 -0800
Subject: [PATCH] Let text_utils work with strings that contain ansi escape
 sequences.

---
 string_utils.py | 24 ++++++++++++++++++++----
 text_utils.py   | 32 ++++++++++++++++++--------------
 2 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/string_utils.py b/string_utils.py
index a766ada..37851e0 100644
--- a/string_utils.py
+++ b/string_utils.py
@@ -1218,6 +1218,22 @@ def sprintf(*args, **kwargs) -> str:
     return ret
 
 
+def strip_ansi_sequences(in_str: str) -> str:
+    """Strips ANSI sequences out of strings.
+
+    >>> import ansi as a
+    >>> s = a.fg('blue') + 'blue!' + a.reset()
+    >>> len(s)   # '\x1b[38;5;21mblue!\x1b[m'
+    18
+    >>> len(strip_ansi_sequences(s))
+    5
+    >>> strip_ansi_sequences(s)
+    'blue!'
+
+    """
+    return re.sub(r'\x1b\[[\d+;]*[a-z]', '', in_str)
+
+
 class SprintfStdout(contextlib.AbstractContextManager):
     """
     A context manager that captures outputs to stdout.
@@ -1680,16 +1696,16 @@ def replace_all(in_str: str, replace_set: str, replacement: str) -> str:
     return in_str
 
 
-def replace_nth(string: str, source: str, target: str, nth: int):
+def replace_nth(in_str: str, source: str, target: str, nth: int):
     """Replaces the nth occurrance of a substring within a string.
 
     >>> replace_nth('this is a test', ' ', '-', 3)
     'this is a-test'
 
     """
-    where = [m.start() for m in re.finditer(source, string)][nth - 1]
-    before = string[:where]
-    after = string[where:]
+    where = [m.start() for m in re.finditer(source, in_str)][nth - 1]
+    before = in_str[:where]
+    after = in_str[where:]
     after = after.replace(source, target, 1)
     return before + after
 
diff --git a/text_utils.py b/text_utils.py
index 55b0575..cc1269c 100644
--- a/text_utils.py
+++ b/text_utils.py
@@ -12,6 +12,7 @@ from collections import defaultdict
 from dataclasses import dataclass
 from typing import Dict, Generator, List, Literal, Optional, Tuple
 
+import string_utils
 from ansi import fg, reset
 
 logger = logging.getLogger(__file__)
@@ -148,9 +149,9 @@ def distribute_strings(
 
     """
     ret = ' ' + ' '.join(strings) + ' '
-    assert len(ret) < width
+    assert len(string_utils.strip_ansi_sequences(ret)) < width
     x = 0
-    while len(ret) < width:
+    while len(string_utils.strip_ansi_sequences(ret)) < width:
         spaces = [m.start() for m in re.finditer(r' ([^ ]|$)', ret)]
         where = spaces[x]
         before = ret[:where]
@@ -172,10 +173,12 @@ def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") ->
     'This   is   a   test'
 
     """
-    assert len(string) <= width
+    assert len(string_utils.strip_ansi_sequences(string)) <= width
     padding = padding[0]
     first, *rest, last = string.split()
-    w = width - (len(first) + len(last))
+    w = width - (
+        len(string_utils.strip_ansi_sequences(first)) + len(string_utils.strip_ansi_sequences(last))
+    )
     ret = first + distribute_strings(rest, width=w, padding=padding) + last
     return ret
 
@@ -197,7 +200,7 @@ def justify_string(
     """
     alignment = alignment[0]
     padding = padding[0]
-    while len(string) < width:
+    while len(string_utils.strip_ansi_sequences(string)) < width:
         if alignment == "l":
             string += padding
         elif alignment == "r":
@@ -226,13 +229,16 @@ def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
     retval = ""
     line = ""
     for word in text.split():
-        if len(line) + len(word) > width:
+        if (
+            len(string_utils.strip_ansi_sequences(line))
+            + len(string_utils.strip_ansi_sequences(word))
+        ) > width:
             line = line[1:]
             line = justify_string(line, width=width, alignment=alignment)
             retval = retval + "\n" + line
             line = ""
         line = line + " " + word
-    if len(line) > 0:
+    if len(string_utils.strip_ansi_sequences(line)) > 0:
         retval += "\n" + line[1:]
     return retval[1:]
 
@@ -241,7 +247,7 @@ def generate_padded_columns(text: List[str]) -> Generator:
     max_width: Dict[int, int] = defaultdict(int)
     for line in text:
         for pos, word in enumerate(line.split()):
-            max_width[pos] = max(max_width[pos], len(word))
+            max_width[pos] = max(max_width[pos], len(string_utils.strip_ansi_sequences(word)))
 
     for line in text:
         out = ""
@@ -257,11 +263,11 @@ def wrap_string(text: str, n: int) -> str:
     out = ''
     width = 0
     for chunk in chunks:
-        if width + len(chunk) > n:
+        if width + len(string_utils.strip_ansi_sequences(chunk)) > n:
             out += '\n'
             width = 0
         out += chunk + ' '
-        width += len(chunk) + 1
+        width += len(string_utils.strip_ansi_sequences(chunk)) + 1
     return out
 
 
@@ -301,8 +307,6 @@ class Indenter(contextlib.AbstractContextManager):
         return False
 
     def print(self, *arg, **kwargs):
-        import string_utils
-
         text = string_utils.sprintf(*arg, **kwargs)
         print(self.pad_prefix + self.padding * self.level + text, end='')
 
@@ -316,7 +320,7 @@ def header(title: str, *, width: int = 80, color: str = ''):
 
     """
     w = width
-    w -= len(title) + 4
+    w -= len(string_utils.strip_ansi_sequences(title)) + 4
     if w >= 4:
         left = 4 * '-'
         right = (w - 4) * '-'
@@ -354,7 +358,7 @@ def box(
         ret += color + '│' + ' ' * w + '│' + rset + '\n'
     if text is not None:
         for line in justify_text(text, width=w - 2, alignment='l').split('\n'):
-            tw = len(line)
+            tw = len(string_utils.strip_ansi_sequences(line))
             assert tw < w
             ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
     ret += color + '╰' + '─' * w + '╯' + rset + '\n'
-- 
2.47.1