Try multiple strategies to determine the console size.

[python_utils.git] / text_utils.py

diff --git a/text_utils.py b/text_utils.py
index afe0f63dcf1f86674fc95ba9d8dcca26747b158c..46f3756eadea047db2b2c311ccca5ffe442dad15 100644 (file)
--- a/text_utils.py
+++ b/text_utils.py
@@ -1,11 +1,14 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+# © Copyright 2021-2022, Scott Gasch
+
 """Utilities for dealing with "text"."""
 
 import contextlib
 import logging
 import math
+import os
 import re
 import sys
 from collections import defaultdict
@@ -31,15 +34,39 @@ def get_console_rows_columns() -> RowsColumns:
 
     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:
+        try:
+            rows, cols = cmd(
+                "stty size",
+                timeout_seconds=1.0,
+            ).split()
+        except Exception:
+            rows = None
+            cols = None
+
+    if rows is None:
+        try:
+            rows = cmd(
+                "tput rows",
+                timeout_seconds=1.0,
+            )
+        except Exception:
+            rows = None
+
+    if cols is None:
+        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))
 
 
 def progress_graph(
@@ -324,6 +351,7 @@ 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.
@@ -333,19 +361,27 @@ def header(
 
     """
     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,11 +395,26 @@ 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:
+    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:
     assert width > 4
     ret = ''
@@ -386,9 +437,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
@@ -412,7 +463,7 @@ def print_box(
     ╰────╯
 
     """
-    print(box(title, text, width=width, color=color), end='')
+    print(preformatted_box(title, text, width=width, color=color), end='')
 
 
 if __name__ == '__main__':