Try multiple strategies to determine the console size.
[python_utils.git] / text_utils.py
index 55b057569b531f211e0073c0dc931ae4fd239bd4..46f3756eadea047db2b2c311ccca5ffe442dad15 100644 (file)
@@ -1,17 +1,21 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+# © Copyright 2021-2022, Scott Gasch
+
 """Utilities for dealing with "text"."""
 
 import contextlib
 import logging
 import math
 """Utilities for dealing with "text"."""
 
 import contextlib
 import logging
 import math
+import os
 import re
 import sys
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import Dict, Generator, List, Literal, Optional, Tuple
 
 import re
 import sys
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import Dict, Generator, List, Literal, Optional, Tuple
 
+import string_utils
 from ansi import fg, reset
 
 logger = logging.getLogger(__file__)
 from ansi import fg, reset
 
 logger = logging.getLogger(__file__)
@@ -30,15 +34,39 @@ def get_console_rows_columns() -> RowsColumns:
 
     from exec_utils import cmd
 
 
     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(
 
 
 def progress_graph(
@@ -148,9 +176,9 @@ def distribute_strings(
 
     """
     ret = ' ' + ' '.join(strings) + ' '
 
     """
     ret = ' ' + ' '.join(strings) + ' '
-    assert len(ret) < width
+    assert len(string_utils.strip_ansi_sequences(ret)) < width
     x = 0
     x = 0
-    while len(ret) < width:
+    while len(string_utils.strip_ansi_sequences(ret)) < width:
         spaces = [m.start() for m in re.finditer(r' ([^ ]|$)', ret)]
         where = spaces[x]
         before = ret[:where]
         spaces = [m.start() for m in re.finditer(r' ([^ ]|$)', ret)]
         where = spaces[x]
         before = ret[:where]
@@ -172,10 +200,12 @@ def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") ->
     'This   is   a   test'
 
     """
     'This   is   a   test'
 
     """
-    assert len(string) <= width
+    assert len(string_utils.strip_ansi_sequences(string)) <= width
     padding = padding[0]
     first, *rest, last = string.split()
     padding = padding[0]
     first, *rest, last = string.split()
-    w = width - (len(first) + len(last))
+    w = width - (
+        len(string_utils.strip_ansi_sequences(first)) + len(string_utils.strip_ansi_sequences(last))
+    )
     ret = first + distribute_strings(rest, width=w, padding=padding) + last
     return ret
 
     ret = first + distribute_strings(rest, width=w, padding=padding) + last
     return ret
 
@@ -197,7 +227,7 @@ def justify_string(
     """
     alignment = alignment[0]
     padding = padding[0]
     """
     alignment = alignment[0]
     padding = padding[0]
-    while len(string) < width:
+    while len(string_utils.strip_ansi_sequences(string)) < width:
         if alignment == "l":
             string += padding
         elif alignment == "r":
         if alignment == "l":
             string += padding
         elif alignment == "r":
@@ -214,26 +244,36 @@ def justify_string(
     return string
 
 
     return string
 
 
-def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
+def justify_text(text: str, *, width: int = 80, alignment: str = "c", indent_by: int = 0) -> str:
     """
     """
-    Justifies text.
+    Justifies text optionally with initial indentation.
 
     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
 
     """
 
     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
 
     """
-    retval = ""
-    line = ""
+    retval = ''
+    indent = ''
+    if indent_by > 0:
+        indent += ' ' * indent_by
+    line = indent
+
     for word in text.split():
     for word in text.split():
-        if len(line) + len(word) > width:
+        if (
+            len(string_utils.strip_ansi_sequences(line))
+            + len(string_utils.strip_ansi_sequences(word))
+        ) > width:
             line = line[1:]
             line = justify_string(line, width=width, alignment=alignment)
             line = line[1:]
             line = justify_string(line, width=width, alignment=alignment)
-            retval = retval + "\n" + line
-            line = ""
-        line = line + " " + word
-    if len(line) > 0:
-        retval += "\n" + line[1:]
+            retval = retval + '\n' + line
+            line = indent
+        line = line + ' ' + word
+    if len(string_utils.strip_ansi_sequences(line)) > 0:
+        if alignment != 'j':
+            retval += "\n" + justify_string(line[1:], width=width, alignment=alignment)
+        else:
+            retval += "\n" + line[1:]
     return retval[1:]
 
 
     return retval[1:]
 
 
@@ -241,7 +281,7 @@ 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: 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))
+            max_width[pos] = max(max_width[pos], len(string_utils.strip_ansi_sequences(word)))
 
     for line in text:
         out = ""
 
     for line in text:
         out = ""
@@ -257,11 +297,11 @@ def wrap_string(text: str, n: int) -> str:
     out = ''
     width = 0
     for chunk in chunks:
     out = ''
     width = 0
     for chunk in chunks:
-        if width + len(chunk) > n:
+        if width + len(string_utils.strip_ansi_sequences(chunk)) > n:
             out += '\n'
             width = 0
         out += chunk + ' '
             out += '\n'
             width = 0
         out += chunk + ' '
-        width += len(chunk) + 1
+        width += len(string_utils.strip_ansi_sequences(chunk)) + 1
     return out
 
 
     return out
 
 
@@ -301,36 +341,80 @@ class Indenter(contextlib.AbstractContextManager):
         return False
 
     def print(self, *arg, **kwargs):
         return False
 
     def print(self, *arg, **kwargs):
-        import string_utils
-
         text = string_utils.sprintf(*arg, **kwargs)
         print(self.pad_prefix + self.padding * self.level + text, end='')
 
 
         text = string_utils.sprintf(*arg, **kwargs)
         print(self.pad_prefix + self.padding * self.level + text, end='')
 
 
-def header(title: str, *, width: int = 80, color: str = ''):
+def header(
+    title: str,
+    *,
+    width: Optional[int] = None,
+    align: Optional[str] = None,
+    style: Optional[str] = 'solid',
+    color: Optional[str] = None,
+):
     """
     Returns a nice header line with a title.
 
     """
     Returns a nice header line with a title.
 
-    >>> header('title', width=60, color='')
+    >>> header('title', width=60, style='ascii')
     '----[ title ]-----------------------------------------------'
 
     """
     '----[ title ]-----------------------------------------------'
 
     """
-    w = width
-    w -= len(title) + 4
-    if w >= 4:
-        left = 4 * '-'
-        right = (w - 4) * '-'
-        if color != '' and color is not None:
-            r = reset()
-        else:
-            r = ''
-        return f'{left}[ {color}{title}{r} ]{right}'
+    if not width:
+        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 + text_len + 4)
+    elif align == 'right':
+        right = 4
+        left = width - (right + text_len + 4)
     else:
     else:
-        return ''
+        left = int((width - (text_len + 4)) / 2)
+        right = left
+        while left + text_len + 4 + right < width:
+            right += 1
+
+    if style == 'solid':
+        line_char = '━'
+        begin = ''
+        end = ''
+    elif style == 'dashed':
+        line_char = '┅'
+        begin = ''
+        end = ''
+    else:
+        line_char = '-'
+        begin = '['
+        end = ']'
+    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 = ''
 
 
 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 = ''
 ) -> str:
     assert width > 4
     ret = ''
@@ -353,9 +437,9 @@ def box(
         )
         ret += color + '│' + ' ' * w + '│' + rset + '\n'
     if text is not None:
         )
         ret += color + '│' + ' ' * w + '│' + rset + '\n'
     if text is not None:
-        for line in justify_text(text, width=w - 2, alignment='l').split('\n'):
-            tw = len(line)
-            assert tw < w
+        for line in text.split('\n'):
+            tw = len(string_utils.strip_ansi_sequences(line))
+            assert tw <= w
             ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
     ret += color + '╰' + '─' * w + '╯' + rset + '\n'
     return ret
             ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
     ret += color + '╰' + '─' * w + '╯' + rset + '\n'
     return ret
@@ -379,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__':
 
 
 if __name__ == '__main__':