Easier and more self documenting patterns for loading/saving Persistent
[python_utils.git] / text_utils.py
index afe0f63dcf1f86674fc95ba9d8dcca26747b158c..39b8fe3e3db64266aeaeaad4d92503be393017e9 100644 (file)
@@ -1,11 +1,15 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+# © Copyright 2021-2022, Scott Gasch
+
 """Utilities for dealing with "text"."""
 
 import contextlib
+import enum
 import logging
 import math
+import os
 import re
 import sys
 from collections import defaultdict
@@ -23,42 +27,101 @@ class RowsColumns:
     """Row + Column"""
 
     rows: int = 0
+    """Numer of rows"""
+
     columns: int = 0
+    """Number of columns"""
 
 
 def get_console_rows_columns() -> RowsColumns:
-    """Returns the number of rows/columns on the current console."""
-
+    """
+    Returns:
+        The number of rows/columns on the current console or None
+        if we can't tell or an error occurred.
+    """
     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?!') from e
-    return RowsColumns(int(rows), int(columns))
+    rows: Optional[str] = os.environ.get('LINES', None)
+    cols: Optional[str] = os.environ.get('COLUMNS', None)
+    if not rows or not cols:
+        logger.debug('Rows: %s, cols: %s, trying stty.', rows, cols)
+        try:
+            rows, cols = cmd(
+                "stty size",
+                timeout_seconds=1.0,
+            ).split()
+        except Exception:
+            rows = None
+            cols = None
+
+    if rows is None:
+        logger.debug('Rows: %s, cols: %s, tput rows.', rows, cols)
+        try:
+            rows = cmd(
+                "tput rows",
+                timeout_seconds=1.0,
+            )
+        except Exception:
+            rows = None
+
+    if cols is None:
+        logger.debug('Rows: %s, cols: %s, tput cols.', rows, cols)
+        try:
+            cols = cmd(
+                "tput cols",
+                timeout_seconds=1.0,
+            )
+        except Exception:
+            cols = None
+
+    if not rows or not cols:
+        raise Exception('Can\'t determine console size?!')
+    return RowsColumns(int(rows), int(cols))
+
+
+class BarGraphText(enum.Enum):
+    """What kind of text to include at the end of the bar graph?"""
+
+    NONE = (0,)
+    """None, leave it blank."""
+
+    PERCENTAGE = (1,)
+    """XX.X%"""
+
+    FRACTION = (2,)
+    """N / K"""
 
 
-def progress_graph(
+def bar_graph(
     current: int,
     total: int,
     *,
     width=70,
+    text: BarGraphText = BarGraphText.PERCENTAGE,
     fgcolor=fg("school bus yellow"),
     left_end="[",
     right_end="]",
     redraw=True,
 ) -> None:
-    """Draws a progress graph."""
-
-    percent = current / total
+    """Draws a progress graph at the current cursor position.
+
+    Args:
+        current: how many have we done so far?
+        total: how many are there to do total?
+        text: how should we render the text at the end?
+        width: how many columns wide should be progress graph be?
+        fgcolor: what color should "done" part of the graph be?
+        left_end: the character at the left side of the graph
+        right_end: the character at the right side of the graph
+        redraw: if True, omit a line feed after the carriage return
+            so that subsequent calls to this method redraw the graph
+            iteratively.
+    """
     ret = "\r" if redraw else "\n"
-    bar = bar_graph(
-        percent,
-        include_text=True,
+    bar = bar_graph_string(
+        current,
+        total,
+        text=text,
         width=width,
         fgcolor=fgcolor,
         left_end=left_end,
@@ -67,10 +130,21 @@ def progress_graph(
     print(bar, end=ret, flush=True, file=sys.stderr)
 
 
-def bar_graph(
-    percentage: float,
+def _make_bar_graph_text(text: BarGraphText, current: int, total: int, percentage: float):
+    if text == BarGraphText.NONE:
+        return ""
+    elif text == BarGraphText.PERCENTAGE:
+        return f'{percentage:.1f}'
+    elif text == BarGraphText.FRACTION:
+        return f'{current} / {total}'
+    raise ValueError(text)
+
+
+def bar_graph_string(
+    current: int,
+    total: int,
     *,
-    include_text=True,
+    text: BarGraphText = BarGraphText.PERCENTAGE,
     width=70,
     fgcolor=fg("school bus yellow"),
     reset_seq=reset(),
@@ -79,17 +153,28 @@ def bar_graph(
 ) -> str:
     """Returns a string containing a bar graph.
 
-    >>> bar_graph(0.5, fgcolor='', reset_seq='')
+    Args:
+        current: how many have we done so far?
+        total: how many are there to do total?
+        text: how should we render the text at the end?
+        width: how many columns wide should be progress graph be?
+        fgcolor: what color should "done" part of the graph be?
+        reset_seq: sequence to use to turn off color
+        left_end: the character at the left side of the graph
+        right_end: the character at the right side of the graph
+
+    >>> bar_graph(5, 10, fgcolor='', reset_seq='')
     '[███████████████████████████████████                                   ] 50.0%'
 
     """
 
+    if total != 0:
+        percentage = float(current) / float(total)
+    else:
+        percentage = 0.0
     if percentage < 0.0 or percentage > 1.0:
         raise ValueError(percentage)
-    if include_text:
-        text = f"{percentage*100.0:2.1f}%"
-    else:
-        text = ""
+    text = _make_bar_graph_text(text, current, total, percentage)
     whole_width = math.floor(percentage * width)
     if whole_width == width:
         whole_width -= 1
@@ -117,6 +202,16 @@ def sparkline(numbers: List[float]) -> Tuple[float, float, str]:
     """
     Makes a "sparkline" little inline histogram graph.  Auto scales.
 
+    Args:
+        numbers: the population over which to create the sparkline
+
+    Returns:
+        a three tuple containing:
+
+        * the minimum number in the population
+        * the maximum number in the population
+        * a string representation of the population in a concise format
+
     >>> sparkline([1, 2, 3, 5, 10, 3, 5, 7])
     (1, 10, '▁▁▂▄█▂▄▆')
 
@@ -144,9 +239,16 @@ def distribute_strings(
     """
     Distributes strings into a line for justified text.
 
+    Args:
+        strings: a list of string tokens to distribute
+        width: the width of the line to create
+        padding: the padding character to place between string chunks
+
+    Returns:
+        The distributed, justified string.
+
     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
     '      this      is      a      test     '
-
     """
     ret = ' ' + ' '.join(strings) + ' '
     assert len(string_utils.strip_ansi_sequences(ret)) < width
@@ -163,13 +265,21 @@ def distribute_strings(
     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:
     """
-    Justifies a string.
+    Justifies a string chunk by chunk.
+
+    Args:
+        string: the string to be justified
+        width: how wide to make the output
+        padding: what padding character to use between chunks
+
+    Returns:
+        the justified string
 
-    >>> justify_string_by_chunk("This is a test", 40)
+    >>> _justify_string_by_chunk("This is a test", 40)
     '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'
 
     """
@@ -186,7 +296,18 @@ def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") ->
 def justify_string(
     string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
 ) -> str:
-    """Justify a string.
+    """Justify a string to width with left, right, center of justified
+    alignment.
+
+    Args:
+        string: the string to justify
+        width: the width to justify the string to
+        alignment: a single character indicating the desired alignment:
+            * 'c' = centered within the width
+            * 'j' = justified at width
+            * 'l' = left alignment
+            * 'r' = right alignment
+        padding: the padding character to use while justifying
 
     >>> justify_string('This is another test', width=40, alignment='c')
     '          This is another test          '
@@ -196,7 +317,6 @@ def justify_string(
     '                    This is another test'
     >>> justify_string('This is another test', width=40, alignment='j')
     'This        is        another       test'
-
     """
     alignment = alignment[0]
     padding = padding[0]
@@ -206,7 +326,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
@@ -218,8 +338,21 @@ def justify_string(
 
 
 def justify_text(text: str, *, width: int = 80, alignment: str = "c", indent_by: int = 0) -> str:
-    """
-    Justifies text optionally with initial indentation.
+    """Justifies text with left, right, centered or justified alignment
+    and optionally with initial indentation.
+
+    Args:
+        text: the text to be justified
+        width: the width at which to justify text
+        alignment: a single character indicating the desired alignment:
+            * 'c' = centered within the width
+            * 'j' = justified at width
+            * 'l' = left alignment
+            * 'r' = right alignment
+        indent_by: if non-zero, adds n prefix spaces to indent the text.
+
+    Returns:
+        The justified text.
 
     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
@@ -251,6 +384,26 @@ def justify_text(text: str, *, width: int = 80, alignment: str = "c", indent_by:
 
 
 def generate_padded_columns(text: List[str]) -> Generator:
+    """Given a list of strings, break them into columns using :meth:`split`
+    and then compute the maximum width of each column.  Finally,
+    distribute the columular chunks into the output padding each to
+    the proper width.
+
+    Args:
+        text: a list of strings to chunk into padded columns
+
+    Returns:
+        padded columns based on text.split()
+
+    >>> for x in generate_padded_columns(
+    ...     [ 'reading writing arithmetic',
+    ...       'mathematics psychology physics',
+    ...       'communications sociology anthropology' ]):
+    ...     print(x.strip())
+    reading        writing    arithmetic
+    mathematics    psychology physics
+    communications sociology  anthropology
+    """
     max_width: Dict[int, int] = defaultdict(int)
     for line in text:
         for pos, word in enumerate(line.split()):
@@ -266,6 +419,14 @@ def generate_padded_columns(text: List[str]) -> Generator:
 
 
 def wrap_string(text: str, n: int) -> str:
+    """
+    Args:
+        text: the string to be wrapped
+        n: the width after which to wrap text
+
+    Returns:
+        The wrapped form of text
+    """
     chunks = text.split()
     out = ''
     width = 0
@@ -280,13 +441,20 @@ def wrap_string(text: str, n: int) -> str:
 
 class Indenter(contextlib.AbstractContextManager):
     """
-    with Indenter(pad_count = 8) as i:
-        i.print('test')
-        with i:
-            i.print('-ing')
+    Context manager that indents stuff (even recursively).  e.g.::
+
+        with Indenter(pad_count = 8) as i:
+            i.print('test')
             with i:
-                i.print('1, 2, 3')
+                i.print('-ing')
+                with i:
+                    i.print('1, 2, 3')
+
+    Yields::
 
+        test
+                -ing
+                        1, 2, 3
     """
 
     def __init__(
@@ -296,6 +464,13 @@ class Indenter(contextlib.AbstractContextManager):
         pad_char: str = ' ',
         pad_count: int = 4,
     ):
+        """Construct an Indenter.
+
+        Args:
+            pad_prefix: an optional prefix to prepend to each line
+            pad_char: the character used to indent
+            pad_count: the number of pad_chars to use to indent
+        """
         self.level = -1
         if pad_prefix is not None:
             self.pad_prefix = pad_prefix
@@ -324,28 +499,45 @@ def header(
     width: Optional[int] = None,
     align: Optional[str] = None,
     style: Optional[str] = 'solid',
+    color: Optional[str] = None,
 ):
     """
-    Returns a nice header line with a title.
+    Creates a nice header line with a title.
+
+    Args:
+        title: the title
+        width: how wide to make the header
+        align: "left" or "right"
+        style: "ascii", "solid" or "dashed"
+
+    Returns:
+        The header as a string.
 
     >>> header('title', width=60, style='ascii')
     '----[ title ]-----------------------------------------------'
-
     """
     if not width:
-        width = get_console_rows_columns().columns
+        try:
+            width = get_console_rows_columns().columns
+        except Exception:
+            width = 80
     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 + len(string_utils.strip_ansi_sequences(title)) + 4)
+        right = width - (left + text_len + 4)
     elif align == 'right':
         right = 4
-        left = width - (right + len(string_utils.strip_ansi_sequences(title)) + 4)
+        left = width - (right + text_len + 4)
     else:
-        left = int((width - (len(string_utils.strip_ansi_sequences(title)) + 4)) / 2)
+        left = int((width - (text_len + 4)) / 2)
         right = left
+        while left + text_len + 4 + right < width:
+            right += 1
 
     if style == 'solid':
         line_char = '━'
@@ -359,12 +551,68 @@ def header(
         line_char = '-'
         begin = '['
         end = ']'
-    return line_char * left + begin + ' ' + title + ' ' + end + line_char * right
+    if color:
+        col = color
+        reset_seq = reset()
+    else:
+        col = ''
+        reset_seq = ''
+    return line_char * left + begin + col + ' ' + title + ' ' + reset_seq + end + line_char * right
 
 
 def box(
     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
 ) -> str:
+    """
+    Make a nice unicode box (optionally with color) around some text.
+
+    Args:
+        title: the title of the box
+        text: the text in the box
+        width: the box's width
+        color: the box's color
+
+    Returns:
+        the box as a string
+
+    >>> print(box('title', 'this is some text', width=20).strip())
+    ╭──────────────────╮
+    │       title      │
+    │                  │
+    │ this is some     │
+    │ text             │
+    ╰──────────────────╯
+    """
+    assert width > 4
+    if text is not None:
+        text = justify_text(text, width=width - 4, alignment='l')
+    return preformatted_box(title, text, width=width, color=color)
+
+
+def preformatted_box(
+    title: Optional[str] = None, text: Optional[str] = None, *, width=80, color: str = ''
+) -> str:
+    """Creates a nice box with rounded corners and returns it as a string.
+
+    Args:
+        title: the title of the box
+        text: the text inside the box
+        width: the width of the box
+        color: the box's color
+
+    Returns:
+        the box as a string
+
+    >>> print(preformatted_box('title', 'this\\nis\\nsome\\ntext', width=20).strip())
+    ╭──────────────────╮
+    │       title      │
+    │                  │
+    │ this             │
+    │ is               │
+    │ some             │
+    │ text             │
+    ╰──────────────────╯
+    """
     assert width > 4
     ret = ''
     if color == '':
@@ -386,9 +634,9 @@ def box(
         )
         ret += color + '│' + ' ' * w + '│' + rset + '\n'
     if text is not None:
-        for line in justify_text(text, width=w - 2, alignment='l').split('\n'):
+        for line in text.split('\n'):
             tw = len(string_utils.strip_ansi_sequences(line))
-            assert tw < w
+            assert tw <= w
             ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
     ret += color + '╰' + '─' * w + '╯' + rset + '\n'
     return ret
@@ -410,9 +658,8 @@ def print_box(
     ╭────╮
     │ OK │
     ╰────╯
-
     """
-    print(box(title, text, width=width, color=color), end='')
+    print(preformatted_box(title, text, width=width, color=color), end='')
 
 
 if __name__ == '__main__':