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 a6337d6c5382b80ee75a7c199849123a7e97cea8..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,7 +361,10 @@ 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:
@@ -364,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 = ''
@@ -391,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
@@ -417,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__':