Let text_utils work with strings that contain ansi escape sequences.

[python_utils.git] / text_utils.py

diff --git a/text_utils.py b/text_utils.py
index 564d67eb98c479efdfa00671a98ea179f0f189a8..cc1269c58f292039cae5ecb93b533905020360f9 100644 (file)
--- a/text_utils.py
+++ b/text_utils.py
@@ -6,11 +6,13 @@
 import contextlib
 import logging
 import math
 import contextlib
 import logging
 import math

+import re

 import sys
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import Dict, Generator, List, Literal, Optional, Tuple
 
 import sys
 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__)
 from ansi import fg, reset
 
 logger = logging.getLogger(__file__)


@@ -137,30 +139,28 @@ def distribute_strings(
     strings: List[str],
     *,
     width: int = 80,
     strings: List[str],
     *,
     width: int = 80,

-    alignment: str = "c",

     padding: str = " ",
 ) -> str:
     """
     padding: str = " ",
 ) -> str:
     """

-    Distributes strings into a line with a particular justification.
+    Distributes strings into a line for justified text.

 
     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
 
     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)

-    '   this       is         a       test   '
-    >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='l')
-    'this      is        a         test      '
-    >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='r')
-    '      this        is         a      test'
+    '      this      is      a      test     '

 
     """
 
     """

-    subwidth = math.floor(width / len(strings))
-    retval = ""
-    for string in strings:
-        string = justify_string(string, width=subwidth, alignment=alignment, padding=padding)
-        retval += string
-    while len(retval) > width:
-        retval = retval.replace('  ', ' ', 1)
-    while len(retval) < width:
-        retval = retval.replace(' ', '  ', 1)
-    return retval
+    ret = ' ' + ' '.join(strings) + ' '
+    assert len(string_utils.strip_ansi_sequences(ret)) < width
+    x = 0
+    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]
+        after = ret[where:]
+        ret = before + padding + after
+        x += 1
+        if x >= len(spaces):
+            x = 0
+    return ret

 
 
 def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:
 
 
 def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:


@@ -168,15 +168,18 @@ def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") ->
     Justifies a string.
 
     >>> justify_string_by_chunk("This is a test", 40)
     Justifies a string.
 
     >>> justify_string_by_chunk("This is a test", 40)

-    'This       is              a        test'
+    'This          is          a         test'

     >>> justify_string_by_chunk("This is a test", 20)
     >>> justify_string_by_chunk("This is a test", 20)

-    'This  is    a   test'
+    'This   is   a   test'

 
     """
 
     """

+    assert len(string_utils.strip_ansi_sequences(string)) <= width

     padding = padding[0]
     first, *rest, last = string.split()
     padding = padding[0]
     first, *rest, last = string.split()

-    w = width - (len(first) + 1 + len(last) + 1)
-    ret = first + padding + distribute_strings(rest, width=w, padding=padding) + padding + 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
 
 
     return ret
 
 


@@ -192,12 +195,12 @@ def justify_string(
     >>> justify_string('This is another test', width=40, alignment='r')
     '                    This is another test'
     >>> justify_string('This is another test', width=40, alignment='j')
     >>> justify_string('This is another test', width=40, alignment='r')
     '                    This is another test'
     >>> justify_string('This is another test', width=40, alignment='j')

-    'This       is           another     test'
+    'This        is        another       test'

 
     """
     alignment = alignment[0]
     padding = padding[0]
 
     """
     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":
         if alignment == "l":
             string += padding
         elif alignment == "r":


@@ -221,17 +224,21 @@ def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'

+

     """
     retval = ""
     line = ""
     for word in text.split():
     """
     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
             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:]
 
         retval += "\n" + line[1:]
     return retval[1:]
 


@@ -240,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: 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 = ""
 
     for line in text:
         out = ""


@@ -256,11 +263,11 @@ def wrap_string(text: str, n: int) -> str:
     out = ''
     width = 0
     for chunk in chunks:
     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 + ' '
             out += '\n'
             width = 0
         out += chunk + ' '

-        width += len(chunk) + 1
+        width += len(string_utils.strip_ansi_sequences(chunk)) + 1

     return out
 
 
     return out
 
 


@@ -272,6 +279,7 @@ class Indenter(contextlib.AbstractContextManager):
             i.print('-ing')
             with i:
                 i.print('1, 2, 3')
             i.print('-ing')
             with i:
                 i.print('1, 2, 3')

+

     """
 
     def __init__(
     """
 
     def __init__(


@@ -299,8 +307,6 @@ class Indenter(contextlib.AbstractContextManager):
         return False
 
     def print(self, *arg, **kwargs):
         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='')
 
         text = string_utils.sprintf(*arg, **kwargs)
         print(self.pad_prefix + self.padding * self.level + text, end='')
 


@@ -314,7 +320,7 @@ def header(title: str, *, width: int = 80, color: str = ''):
 
     """
     w = width
 
     """
     w = width

-    w -= len(title) + 4
+    w -= len(string_utils.strip_ansi_sequences(title)) + 4

     if w >= 4:
         left = 4 * '-'
         right = (w - 4) * '-'
     if w >= 4:
         left = 4 * '-'
         right = (w - 4) * '-'


@@ -327,6 +333,59 @@ def header(title: str, *, width: int = 80, color: str = ''):
         return ''
 
 
         return ''
 
 

+def box(
+    title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
+) -> str:
+    assert width > 4
+    ret = ''
+    if color == '':
+        rset = ''
+    else:
+        rset = reset()
+    w = width - 2
+    ret += color + '╭' + '─' * w + '╮' + rset + '\n'
+    if title is not None:
+        ret += (
+            color
+            + '│'
+            + rset
+            + justify_string(title, width=w, alignment='c')
+            + color
+            + '│'
+            + rset
+            + '\n'
+        )
+        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(string_utils.strip_ansi_sequences(line))
+            assert tw < w
+            ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
+    ret += color + '╰' + '─' * w + '╯' + rset + '\n'
+    return ret
+
+
+def print_box(
+    title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
+) -> None:
+    """Draws a box with nice rounded corners.
+
+    >>> print_box('Title', 'This is text', width=30)
+    ╭────────────────────────────╮
+    │            Title           │
+    │                            │
+    │ This is text               │
+    ╰────────────────────────────╯
+
+    >>> print_box(None, 'OK', width=6)
+    ╭────╮
+    │ OK │
+    ╰────╯
+
+    """
+    print(box(title, text, width=width, color=color), end='')
+
+

 if __name__ == '__main__':
     import doctest
 
 if __name__ == '__main__':
     import doctest