Used isort to sort imports. Also added to the git pre-commit hook.
[python_utils.git] / text_utils.py
index a71a99dff943df0058cf080100c7feeedf97a6ed..534813c2ef28e186b0adb71ed22b3b1d7cbe4c98 100644 (file)
@@ -2,12 +2,15 @@
 
 """Utilities for dealing with "text"."""
 
+import logging
 import math
 import sys
-from typing import List, NamedTuple
+from collections import defaultdict
+from typing import Dict, Generator, List, NamedTuple, Optional
 
 from ansi import fg, reset
-import exec_utils
+
+logger = logging.getLogger(__file__)
 
 
 class RowsColumns(NamedTuple):
@@ -16,7 +19,18 @@ class RowsColumns(NamedTuple):
 
 
 def get_console_rows_columns() -> RowsColumns:
-    rows, columns = exec_utils.cmd("stty size").split()
+    """Returns the number of rows/columns on the current console."""
+
+    from exec_utils import cmd
+
+    try:
+        rows, columns = cmd(
+            "stty size",
+            timeout_seconds=1.0,
+        ).split()
+    except Exception as e:
+        logger.exception(e)
+        raise Exception('Can\'t determine console size?!')
     return RowsColumns(int(rows), int(columns))
 
 
@@ -30,20 +44,19 @@ def progress_graph(
     right_end="]",
     redraw=True,
 ) -> None:
+    """Draws a progress graph."""
+
     percent = current / total
     ret = "\r" if redraw else "\n"
     bar = bar_graph(
         percent,
-        include_text = True,
-        width = width,
-        fgcolor = fgcolor,
-        left_end = left_end,
-        right_end = right_end)
-    print(
-        bar,
-        end=ret,
-        flush=True,
-        file=sys.stderr)
+        include_text=True,
+        width=width,
+        fgcolor=fgcolor,
+        left_end=left_end,
+        right_end=right_end,
+    )
+    print(bar, end=ret, flush=True, file=sys.stderr)
 
 
 def bar_graph(
@@ -52,9 +65,17 @@ def bar_graph(
     include_text=True,
     width=70,
     fgcolor=fg("school bus yellow"),
+    reset=reset(),
     left_end="[",
     right_end="]",
-) -> None:
+) -> str:
+    """Returns a string containing a bar graph.
+
+    >>> bar_graph(0.5, fgcolor='', reset='')
+    '[███████████████████████████████████                                   ] 50.0%'
+
+    """
+
     if percentage < 0.0 or percentage > 1.0:
         raise ValueError(percentage)
     if include_text:
@@ -65,18 +86,23 @@ def bar_graph(
     if whole_width == width:
         whole_width -= 1
         part_char = "▉"
+    elif whole_width == 0 and percentage > 0.0:
+        part_char = "▏"
     else:
         remainder_width = (percentage * width) % 1
         part_width = math.floor(remainder_width * 8)
         part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width]
     return (
-        left_end +
-        fgcolor +
-        "█" * whole_width + part_char +
-        " " * (width - whole_width - 1) +
-        reset() +
-        right_end + " " +
-        text)
+        left_end
+        + fgcolor
+        + "█" * whole_width
+        + part_char
+        + " " * (width - whole_width - 1)
+        + reset
+        + right_end
+        + " "
+        + text
+    )
 
 
 def distribute_strings(
@@ -86,6 +112,17 @@ def distribute_strings(
     alignment: str = "c",
     padding: str = " ",
 ) -> str:
+    """
+    Distributes strings into a line with a particular justification.
+
+    >>> 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'
+
+    """
     subwidth = math.floor(width / len(strings))
     retval = ""
     for string in strings:
@@ -93,27 +130,51 @@ def distribute_strings(
             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
 
 
-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:
+    """
+    Justifies a string.
+
+    >>> justify_string_by_chunk("This is a test", 40)
+    'This       is              a        test'
+    >>> justify_string_by_chunk("This is a test", 20)
+    'This  is    a   test'
+
+    """
     padding = padding[0]
     first, *rest, last = string.split()
     w = width - (len(first) + 1 + len(last) + 1)
-    retval = (
-        first + padding + distribute_strings(rest, width=w, padding=padding)
+    ret = (
+        first
+        + padding
+        + distribute_strings(rest, width=w, padding=padding)
+        + padding
+        + last
     )
-    while len(retval) + len(last) < width:
-        retval += padding
-    retval += last
-    return retval
+    return ret
 
 
 def justify_string(
     string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
 ) -> str:
+    """Justify a string.
+
+    >>> justify_string('This is another test', width=40, alignment='c')
+    '          This is another test          '
+    >>> justify_string('This is another test', width=40, alignment='l')
+    'This is another test                    '
+    >>> 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'
+
+    """
     alignment = alignment[0]
     padding = padding[0]
     while len(string) < width:
@@ -122,11 +183,7 @@ def justify_string(
         elif alignment == "r":
             string = padding + string
         elif alignment == "j":
-            return justify_string_by_chunk(
-                string,
-                width=width,
-                padding=padding
-            )
+            return justify_string_by_chunk(string, width=width, padding=padding)
         elif alignment == "c":
             if len(string) % 2 == 0:
                 string += padding
@@ -138,7 +195,13 @@ def justify_string(
 
 
 def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
-    print("-" * width)
+    """
+    Justifies text.
+
+    >>> 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():
@@ -151,3 +214,99 @@ def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
     if len(line) > 0:
         retval += "\n" + line[1:]
     return retval[1:]
+
+
+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))
+
+    for line in text:
+        out = ""
+        for pos, word in enumerate(line.split()):
+            width = max_width[pos]
+            word = justify_string(word, width=width, alignment='l')
+            out += f'{word} '
+        yield out
+
+
+def wrap_string(text: str, n: int) -> str:
+    chunks = text.split()
+    out = ''
+    width = 0
+    for chunk in chunks:
+        if width + len(chunk) > n:
+            out += '\n'
+            width = 0
+        out += chunk + ' '
+        width += len(chunk) + 1
+    return out
+
+
+class Indenter(object):
+    """
+    with Indenter(pad_count = 8) as i:
+        i.print('test')
+        with i:
+            i.print('-ing')
+            with i:
+                i.print('1, 2, 3')
+    """
+
+    def __init__(
+        self,
+        *,
+        pad_prefix: Optional[str] = None,
+        pad_char: str = ' ',
+        pad_count: int = 4,
+    ):
+        self.level = -1
+        if pad_prefix is not None:
+            self.pad_prefix = pad_prefix
+        else:
+            self.pad_prefix = ''
+        self.padding = pad_char * pad_count
+
+    def __enter__(self):
+        self.level += 1
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        self.level -= 1
+        if self.level < -1:
+            self.level = -1
+
+    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 = ''):
+    """
+    Returns a nice header line with a title.
+
+    >>> header('title', width=60, color='')
+    '----[ 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}'
+    else:
+        return ''
+
+
+if __name__ == '__main__':
+    import doctest
+
+    doctest.testmod()