Reduce the doctest lease duration...

[python_utils.git] / text_utils.py

diff --git a/text_utils.py b/text_utils.py
index 3b289686d5b4554f1aa1dc40eb639c4b72586702..39b8fe3e3db64266aeaeaad4d92503be393017e9 100644 (file)
--- a/text_utils.py
+++ b/text_utils.py
@@ -6,6 +6,7 @@
 """Utilities for dealing with "text"."""
 
 import contextlib
+import enum
 import logging
 import math
 import os
@@ -43,6 +44,7 @@ def get_console_rows_columns() -> RowsColumns:
     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",
@@ -53,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",
@@ -62,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",
@@ -75,11 +79,25 @@ 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="]",
@@ -90,6 +108,7 @@ def progress_graph(
     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
@@ -98,11 +117,11 @@ def progress_graph(
             so that subsequent calls to this method redraw the graph
             iteratively.
     """
-    percent = current / total
     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,
@@ -111,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(),
@@ -124,25 +154,27 @@ def bar_graph(
     """Returns a string containing a bar graph.
 
     Args:
-        percentage: percentage complete (0..100)
-        include_text: should we include the percentage text at the end?
+        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(0.5, fgcolor='', reset_seq='')
+    >>> 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