Since this thing is on the innerwebs I suppose it should have a
[python_utils.git] / text_utils.py
index 39b694c2721b1c92c73036554a50a6a2194bbc21..f04c61813ef227c4bfe356f42d08c59f2234132c 100644 (file)
@@ -1,6 +1,8 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+# © Copyright 2021-2022, Scott Gasch
+
 """Utilities for dealing with "text"."""
 
 import contextlib
@@ -12,6 +14,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 +151,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 +175,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 +202,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":
@@ -214,25 +219,36 @@ def justify_string(
     return string
 
 
-def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
+def justify_text(text: str, *, width: int = 80, alignment: str = "c", indent_by: int = 0) -> str:
     """
-    Justifies text.
+    Justifies text optionally with initial indentation.
 
     >>> 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 = ""
+    retval = ''
+    indent = ''
+    if indent_by > 0:
+        indent += ' ' * indent_by
+    line = indent
+
     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:
-        retval += "\n" + line[1:]
+            retval = retval + '\n' + line
+            line = indent
+        line = line + ' ' + word
+    if len(string_utils.strip_ansi_sequences(line)) > 0:
+        if alignment != 'j':
+            retval += "\n" + justify_string(line[1:], width=width, alignment=alignment)
+        else:
+            retval += "\n" + line[1:]
     return retval[1:]
 
 
@@ -240,7 +256,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 = ""
@@ -256,11 +272,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
 
 
@@ -272,6 +288,7 @@ class Indenter(contextlib.AbstractContextManager):
             i.print('-ing')
             with i:
                 i.print('1, 2, 3')
+
     """
 
     def __init__(
@@ -299,65 +316,110 @@ 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='')
 
 
-def header(title: str, *, width: int = 80, color: str = ''):
+def header(
+    title: str,
+    *,
+    width: Optional[int] = None,
+    align: Optional[str] = None,
+    style: Optional[str] = 'solid',
+):
     """
     Returns a nice header line with a title.
 
-    >>> header('title', width=60, color='')
+    >>> header('title', width=60, style='ascii')
     '----[ title ]-----------------------------------------------'
 
     """
-    w = width
-    w -= len(title) + 4
-    if w >= 4:
-        left = 4 * '-'
-        right = (w - 4) * '-'
-        if color != '' and color is not None:
-            r = reset()
-        else:
-            r = ''
-        return f'{left}[ {color}{title}{r} ]{right}'
+    if not width:
+        width = get_console_rows_columns().columns
+    if not align:
+        align = 'left'
+    if not style:
+        style = 'ascii'
+
+    text_len = len(string_utils.strip_ansi_sequences(title))
+    if align == 'left':
+        left = 4
+        right = width - (left + text_len + 4)
+    elif align == 'right':
+        right = 4
+        left = width - (right + text_len + 4)
     else:
-        return ''
+        left = int((width - (text_len + 4)) / 2)
+        right = left
+        while left + text_len + 4 + right < width:
+            right += 1
+
+    if style == 'solid':
+        line_char = '━'
+        begin = ''
+        end = ''
+    elif style == 'dashed':
+        line_char = '┅'
+        begin = ''
+        end = ''
+    else:
+        line_char = '-'
+        begin = '['
+        end = ']'
+    return line_char * left + begin + ' ' + title + ' ' + end + line_char * right
 
 
 def box(
     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
-):
-    """Draws a box with nice rounded corners.
-
-    >>> box('Title', 'This is text', width=30)
-    ╭────────────────────────────╮
-    │            Title           │
-    │                            │
-    │ This is text               │
-    ╰────────────────────────────╯
-
-    """
+) -> str:
     assert width > 4
+    ret = ''
     if color == '':
         rset = ''
     else:
         rset = reset()
     w = width - 2
-    print(color + '╭' + '─' * w + '╮' + rset)
+    ret += color + '╭' + '─' * w + '╮' + rset + '\n'
     if title is not None:
-        print(
-            color + '│' + rset + justify_string(title, width=w, alignment='c') + color + '│' + rset
+        ret += (
+            color
+            + '│'
+            + rset
+            + justify_string(title, width=w, alignment='c')
+            + color
+            + '│'
+            + rset
+            + '\n'
         )
-        print(color + '│' + ' ' * w + '│' + rset)
+        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
-            print(color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset)
-    print(color + '╰' + '─' * w + '╯' + rset)
+            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__':