Separate box and print_box.
[python_utils.git] / text_utils.py
index 76cc7e816db9ae87581989bc314fd7d07d4ca154..55b057569b531f211e0073c0dc931ae4fd239bd4 100644 (file)
@@ -3,20 +3,26 @@
 
 """Utilities for dealing with "text"."""
 
+import contextlib
 import logging
 import math
+import re
 import sys
 from collections import defaultdict
-from typing import Dict, Generator, List, NamedTuple, Optional, Tuple
+from dataclasses import dataclass
+from typing import Dict, Generator, List, Literal, Optional, Tuple
 
 from ansi import fg, reset
 
 logger = logging.getLogger(__file__)
 
 
-class RowsColumns(NamedTuple):
-    rows: int
-    columns: int
+@dataclass
+class RowsColumns:
+    """Row + Column"""
+
+    rows: int = 0
+    columns: int = 0
 
 
 def get_console_rows_columns() -> RowsColumns:
@@ -31,7 +37,7 @@ def get_console_rows_columns() -> RowsColumns:
         ).split()
     except Exception as e:
         logger.exception(e)
-        raise Exception('Can\'t determine console size?!')
+        raise Exception('Can\'t determine console size?!') from e
     return RowsColumns(int(rows), int(columns))
 
 
@@ -66,13 +72,13 @@ def bar_graph(
     include_text=True,
     width=70,
     fgcolor=fg("school bus yellow"),
-    reset=reset(),
+    reset_seq=reset(),
     left_end="[",
     right_end="]",
 ) -> str:
     """Returns a string containing a bar graph.
 
-    >>> bar_graph(0.5, fgcolor='', reset='')
+    >>> bar_graph(0.5, fgcolor='', reset_seq='')
     '[███████████████████████████████████                                   ] 50.0%'
 
     """
@@ -99,7 +105,7 @@ def bar_graph(
         + "█" * whole_width
         + part_char
         + " " * (width - whole_width - 1)
-        + reset
+        + reset_seq
         + right_end
         + " "
         + text
@@ -122,40 +128,38 @@ def sparkline(numbers: List[float]) -> Tuple[float, float, str]:
     barcount = len(_bar)
     min_num, max_num = min(numbers), max(numbers)
     span = max_num - min_num
-    sparkline = ''.join(
+    sline = ''.join(
         _bar[min([barcount - 1, int((n - min_num) / span * barcount)])] for n in numbers
     )
-    return min_num, max_num, sparkline
+    return min_num, max_num, sline
 
 
 def distribute_strings(
     strings: List[str],
     *,
     width: int = 80,
-    alignment: str = "c",
     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)
-    '   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(ret) < width
+    x = 0
+    while len(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:
@@ -163,15 +167,16 @@ def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") ->
     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)
-    'This  is    a   test'
+    'This   is   a   test'
 
     """
+    assert len(string) <= width
     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(first) + len(last))
+    ret = first + distribute_strings(rest, width=w, padding=padding) + last
     return ret
 
 
@@ -187,7 +192,7 @@ 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')
-    'This       is           another     test'
+    'This        is        another       test'
 
     """
     alignment = alignment[0]
@@ -216,6 +221,7 @@ 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.'
+
     """
     retval = ""
     line = ""
@@ -259,7 +265,7 @@ def wrap_string(text: str, n: int) -> str:
     return out
 
 
-class Indenter(object):
+class Indenter(contextlib.AbstractContextManager):
     """
     with Indenter(pad_count = 8) as i:
         i.print('test')
@@ -267,6 +273,7 @@ class Indenter(object):
             i.print('-ing')
             with i:
                 i.print('1, 2, 3')
+
     """
 
     def __init__(
@@ -287,10 +294,11 @@ class Indenter(object):
         self.level += 1
         return self
 
-    def __exit__(self, exc_type, exc_value, exc_tb):
+    def __exit__(self, exc_type, exc_value, exc_tb) -> Literal[False]:
         self.level -= 1
         if self.level < -1:
             self.level = -1
+        return False
 
     def print(self, *arg, **kwargs):
         import string_utils
@@ -321,6 +329,59 @@ def header(title: str, *, width: int = 80, color: str = ''):
         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(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