Try multiple strategies to determine the console size.

1 files changed, 46 insertions, 11 deletions

diff --git a/text_utils.py b/text_utils.py
index 18a2257..46f3756 100644
--- a/text_utils.py
+++ b/text_utils.py
@@ -8,6 +8,7 @@
 import contextlib
 import logging
 import math
+import os
 import re
 import sys
 from collections import defaultdict
@@ -33,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(
@@ -326,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.
@@ -335,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:
@@ -366,7 +395,13 @@ 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(