Ran black code formatter on everything.
[python_utils.git] / text_utils.py
index 3be32ff49ec05b2d7ca0978e6cb34b65da64162e..94df3e3499756a86955ebc3ee7330e5329720255 100644 (file)
@@ -3,6 +3,7 @@
 """Utilities for dealing with "text"."""
 
 from collections import defaultdict
+import logging
 import math
 import sys
 from typing import List, NamedTuple, Optional
@@ -10,14 +11,27 @@ from typing import List, NamedTuple, Optional
 from ansi import fg, reset
 
 
+logger = logging.getLogger(__file__)
+
+
 class RowsColumns(NamedTuple):
     rows: int
     columns: int
 
 
 def get_console_rows_columns() -> RowsColumns:
-    from exec_utils import cmd
-    rows, columns = cmd("stty size").split()
+    """Returns the number of rows/columns on the current console."""
+
+    from exec_utils import cmd_with_timeout
+
+    try:
+        rows, columns = cmd_with_timeout(
+            "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))
 
 
@@ -31,20 +45,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(
@@ -53,9 +66,17 @@ def bar_graph(
     include_text=True,
     width=70,
     fgcolor=fg("school bus yellow"),
+    reset=reset(),
     left_end="[",
     right_end="]",
 ) -> None:
+    """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:
@@ -71,13 +92,16 @@ def bar_graph(
         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(
@@ -87,6 +111,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:
@@ -94,27 +129,53 @@ 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:
+    """
+    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:
@@ -123,11 +184,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
@@ -139,7 +196,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():
@@ -182,7 +245,7 @@ def wrap_string(text: str, n: int) -> str:
     return out
 
 
-class Indenter:
+class Indenter(object):
     """
     with Indenter(pad_count = 8) as i:
         i.print('test')
@@ -191,11 +254,14 @@ class Indenter:
             with i:
                 i.print('1, 2, 3')
     """
-    def __init__(self,
-                 *,
-                 pad_prefix: Optional[str] = None,
-                 pad_char: str = ' ',
-                 pad_count: int = 4):
+
+    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
@@ -214,5 +280,34 @@ class Indenter:
 
     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()