Reduce the doctest lease duration...
[python_utils.git] / text_utils.py
index 46f3756eadea047db2b2c311ccca5ffe442dad15..39b8fe3e3db64266aeaeaad4d92503be393017e9 100644 (file)
@@ -6,6 +6,7 @@
 """Utilities for dealing with "text"."""
 
 import contextlib
+import enum
 import logging
 import math
 import os
@@ -26,17 +27,24 @@ 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
 
     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",
@@ -47,6 +55,7 @@ def get_console_rows_columns() -> RowsColumns:
             cols = None
 
     if rows is None:
+        logger.debug('Rows: %s, cols: %s, tput rows.', rows, cols)
         try:
             rows = cmd(
                 "tput rows",
@@ -56,6 +65,7 @@ def get_console_rows_columns() -> RowsColumns:
             rows = None
 
     if cols is None:
+        logger.debug('Rows: %s, cols: %s, tput cols.', rows, cols)
         try:
             cols = cmd(
                 "tput cols",
@@ -69,23 +79,49 @@ def get_console_rows_columns() -> RowsColumns:
     return RowsColumns(int(rows), int(cols))
 
 
-def progress_graph(
+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 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,
@@ -94,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(),
@@ -106,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
@@ -144,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, '▁▁▂▄█▂▄▆')
 
@@ -171,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
@@ -190,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.
 
-    >>> justify_string_by_chunk("This is a test", 40)
+    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)
     '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'
 
     """
@@ -213,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          '
@@ -223,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]
@@ -233,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
@@ -245,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
@@ -278,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()):
@@ -293,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
@@ -307,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__(
@@ -323,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
@@ -354,11 +502,19 @@ def header(
     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:
         try:
@@ -407,6 +563,26 @@ def header(
 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')
@@ -416,6 +592,27 @@ def box(
 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 == '':
@@ -461,7 +658,6 @@ def print_box(
     ╭────╮
     │ OK │
     ╰────╯
-
     """
     print(preformatted_box(title, text, width=width, color=color), end='')