Used isort to sort imports. Also added to the git pre-commit hook.

[python_utils.git] / text_utils.py

diff --git a/text_utils.py b/text_utils.py
index 49ff9b979e2db0e26adfd6dc7b6f0e4663e62289..534813c2ef28e186b0adb71ed22b3b1d7cbe4c98 100644 (file)
--- a/text_utils.py
+++ b/text_utils.py
@@ -2,13 +2,16 @@
 
 """Utilities for dealing with "text"."""
 
-from collections import defaultdict
+import logging
 import math
 import sys
-from typing import List, NamedTuple, Optional
+from collections import defaultdict
+from typing import Dict, Generator, List, NamedTuple, Optional
 
 from ansi import fg, reset
 
+logger = logging.getLogger(__file__)
+
 
 class RowsColumns(NamedTuple):
     rows: int
@@ -19,7 +22,15 @@ def get_console_rows_columns() -> RowsColumns:
     """Returns the number of rows/columns on the current console."""
 
     from exec_utils import cmd
-    rows, columns = cmd("stty size").split()
+
+    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))
 
 
@@ -39,16 +50,13 @@ def progress_graph(
     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(
@@ -60,7 +68,7 @@ def bar_graph(
     reset=reset(),
     left_end="[",
     right_end="]",
-) -> None:
+) -> str:
     """Returns a string containing a bar graph.
 
     >>> bar_graph(0.5, fgcolor='', reset='')
@@ -78,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(
@@ -117,12 +130,14 @@ 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.
 
@@ -135,13 +150,14 @@ def justify_string_by_chunk(
     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(
@@ -167,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
@@ -204,8 +216,8 @@ def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
     return retval[1:]
 
 
-def generate_padded_columns(text: List[str]) -> str:
-    max_width = defaultdict(int)
+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))
@@ -232,7 +244,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')
@@ -241,11 +253,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
@@ -264,6 +279,7 @@ 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='')
 
@@ -277,7 +293,7 @@ def header(title: str, *, width: int = 80, color: str = ''):
 
     """
     w = width
-    w -= (len(title) + 4)
+    w -= len(title) + 4
     if w >= 4:
         left = 4 * '-'
         right = (w - 4) * '-'
@@ -292,4 +308,5 @@ def header(title: str, *, width: int = 80, color: str = ''):
 
 if __name__ == '__main__':
     import doctest
+
     doctest.testmod()