text_utils.py

   1 #!/usr/bin/env python3
   2 # -*- coding: utf-8 -*-
   3
   4 """Utilities for dealing with "text"."""
   5
   6 import contextlib
   7 import logging
   8 import math
   9 import re
  10 import sys
  11 from collections import defaultdict
  12 from dataclasses import dataclass
  13 from typing import Dict, Generator, List, Literal, Optional, Tuple
  14
  15 from ansi import fg, reset
  16
  17 logger = logging.getLogger(__file__)
  18
  19
  20 @dataclass
  21 class RowsColumns:
  22     """Row + Column"""
  23
  24     rows: int = 0
  25     columns: int = 0
  26
  27
  28 def get_console_rows_columns() -> RowsColumns:
  29     """Returns the number of rows/columns on the current console."""
  30
  31     from exec_utils import cmd
  32
  33     try:
  34         rows, columns = cmd(
  35             "stty size",
  36             timeout_seconds=1.0,
  37         ).split()
  38     except Exception as e:
  39         logger.exception(e)
  40         raise Exception('Can\'t determine console size?!') from e
  41     return RowsColumns(int(rows), int(columns))
  42
  43
  44 def progress_graph(
  45     current: int,
  46     total: int,
  47     *,
  48     width=70,
  49     fgcolor=fg("school bus yellow"),
  50     left_end="[",
  51     right_end="]",
  52     redraw=True,
  53 ) -> None:
  54     """Draws a progress graph."""
  55
  56     percent = current / total
  57     ret = "\r" if redraw else "\n"
  58     bar = bar_graph(
  59         percent,
  60         include_text=True,
  61         width=width,
  62         fgcolor=fgcolor,
  63         left_end=left_end,
  64         right_end=right_end,
  65     )
  66     print(bar, end=ret, flush=True, file=sys.stderr)
  67
  68
  69 def bar_graph(
  70     percentage: float,
  71     *,
  72     include_text=True,
  73     width=70,
  74     fgcolor=fg("school bus yellow"),
  75     reset_seq=reset(),
  76     left_end="[",
  77     right_end="]",
  78 ) -> str:
  79     """Returns a string containing a bar graph.
  80
  81     >>> bar_graph(0.5, fgcolor='', reset_seq='')
  82     '[███████████████████████████████████                                   ] 50.0%'
  83
  84     """
  85
  86     if percentage < 0.0 or percentage > 1.0:
  87         raise ValueError(percentage)
  88     if include_text:
  89         text = f"{percentage*100.0:2.1f}%"
  90     else:
  91         text = ""
  92     whole_width = math.floor(percentage * width)
  93     if whole_width == width:
  94         whole_width -= 1
  95         part_char = "▉"
  96     elif whole_width == 0 and percentage > 0.0:
  97         part_char = "▏"
  98     else:
  99         remainder_width = (percentage * width) % 1
 100         part_width = math.floor(remainder_width * 8)
 101         part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width]
 102     return (
 103         left_end
 104         + fgcolor
 105         + "█" * whole_width
 106         + part_char
 107         + " " * (width - whole_width - 1)
 108         + reset_seq
 109         + right_end
 110         + " "
 111         + text
 112     )
 113
 114
 115 def sparkline(numbers: List[float]) -> Tuple[float, float, str]:
 116     """
 117     Makes a "sparkline" little inline histogram graph.  Auto scales.
 118
 119     >>> sparkline([1, 2, 3, 5, 10, 3, 5, 7])
 120     (1, 10, '▁▁▂▄█▂▄▆')
 121
 122     >>> sparkline([104, 99, 93, 96, 82, 77, 85, 73])
 123     (73, 104, '█▇▆▆▃▂▄▁')
 124
 125     """
 126     _bar = '▁▂▃▄▅▆▇█'  # Unicode: 9601, 9602, 9603, 9604, 9605, 9606, 9607, 9608
 127
 128     barcount = len(_bar)
 129     min_num, max_num = min(numbers), max(numbers)
 130     span = max_num - min_num
 131     sline = ''.join(
 132         _bar[min([barcount - 1, int((n - min_num) / span * barcount)])] for n in numbers
 133     )
 134     return min_num, max_num, sline
 135
 136
 137 def distribute_strings(
 138     strings: List[str],
 139     *,
 140     width: int = 80,
 141     padding: str = " ",
 142 ) -> str:
 143     """
 144     Distributes strings into a line for justified text.
 145
 146     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
 147     '      this      is      a      test     '
 148
 149     """
 150     ret = ' ' + ' '.join(strings) + ' '
 151     assert len(ret) < width
 152     x = 0
 153     while len(ret) < width:
 154         spaces = [m.start() for m in re.finditer(r' ([^ ]|$)', ret)]
 155         where = spaces[x]
 156         before = ret[:where]
 157         after = ret[where:]
 158         ret = before + padding + after
 159         x += 1
 160         if x >= len(spaces):
 161             x = 0
 162     return ret
 163
 164
 165 def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:
 166     """
 167     Justifies a string.
 168
 169     >>> justify_string_by_chunk("This is a test", 40)
 170     'This          is          a         test'
 171     >>> justify_string_by_chunk("This is a test", 20)
 172     'This   is   a   test'
 173
 174     """
 175     assert len(string) <= width
 176     padding = padding[0]
 177     first, *rest, last = string.split()
 178     w = width - (len(first) + len(last))
 179     ret = first + distribute_strings(rest, width=w, padding=padding) + last
 180     return ret
 181
 182
 183 def justify_string(
 184     string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
 185 ) -> str:
 186     """Justify a string.
 187
 188     >>> justify_string('This is another test', width=40, alignment='c')
 189     '          This is another test          '
 190     >>> justify_string('This is another test', width=40, alignment='l')
 191     'This is another test                    '
 192     >>> justify_string('This is another test', width=40, alignment='r')
 193     '                    This is another test'
 194     >>> justify_string('This is another test', width=40, alignment='j')
 195     'This        is        another       test'
 196
 197     """
 198     alignment = alignment[0]
 199     padding = padding[0]
 200     while len(string) < width:
 201         if alignment == "l":
 202             string += padding
 203         elif alignment == "r":
 204             string = padding + string
 205         elif alignment == "j":
 206             return justify_string_by_chunk(string, width=width, padding=padding)
 207         elif alignment == "c":
 208             if len(string) % 2 == 0:
 209                 string += padding
 210             else:
 211                 string = padding + string
 212         else:
 213             raise ValueError
 214     return string
 215
 216
 217 def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
 218     """
 219     Justifies text.
 220
 221     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
 222     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
 223     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
 224     """
 225     retval = ""
 226     line = ""
 227     for word in text.split():
 228         if len(line) + len(word) > width:
 229             line = line[1:]
 230             line = justify_string(line, width=width, alignment=alignment)
 231             retval = retval + "\n" + line
 232             line = ""
 233         line = line + " " + word
 234     if len(line) > 0:
 235         retval += "\n" + line[1:]
 236     return retval[1:]
 237
 238
 239 def generate_padded_columns(text: List[str]) -> Generator:
 240     max_width: Dict[int, int] = defaultdict(int)
 241     for line in text:
 242         for pos, word in enumerate(line.split()):
 243             max_width[pos] = max(max_width[pos], len(word))
 244
 245     for line in text:
 246         out = ""
 247         for pos, word in enumerate(line.split()):
 248             width = max_width[pos]
 249             word = justify_string(word, width=width, alignment='l')
 250             out += f'{word} '
 251         yield out
 252
 253
 254 def wrap_string(text: str, n: int) -> str:
 255     chunks = text.split()
 256     out = ''
 257     width = 0
 258     for chunk in chunks:
 259         if width + len(chunk) > n:
 260             out += '\n'
 261             width = 0
 262         out += chunk + ' '
 263         width += len(chunk) + 1
 264     return out
 265
 266
 267 class Indenter(contextlib.AbstractContextManager):
 268     """
 269     with Indenter(pad_count = 8) as i:
 270         i.print('test')
 271         with i:
 272             i.print('-ing')
 273             with i:
 274                 i.print('1, 2, 3')
 275     """
 276
 277     def __init__(
 278         self,
 279         *,
 280         pad_prefix: Optional[str] = None,
 281         pad_char: str = ' ',
 282         pad_count: int = 4,
 283     ):
 284         self.level = -1
 285         if pad_prefix is not None:
 286             self.pad_prefix = pad_prefix
 287         else:
 288             self.pad_prefix = ''
 289         self.padding = pad_char * pad_count
 290
 291     def __enter__(self):
 292         self.level += 1
 293         return self
 294
 295     def __exit__(self, exc_type, exc_value, exc_tb) -> Literal[False]:
 296         self.level -= 1
 297         if self.level < -1:
 298             self.level = -1
 299         return False
 300
 301     def print(self, *arg, **kwargs):
 302         import string_utils
 303
 304         text = string_utils.sprintf(*arg, **kwargs)
 305         print(self.pad_prefix + self.padding * self.level + text, end='')
 306
 307
 308 def header(title: str, *, width: int = 80, color: str = ''):
 309     """
 310     Returns a nice header line with a title.
 311
 312     >>> header('title', width=60, color='')
 313     '----[ title ]-----------------------------------------------'
 314
 315     """
 316     w = width
 317     w -= len(title) + 4
 318     if w >= 4:
 319         left = 4 * '-'
 320         right = (w - 4) * '-'
 321         if color != '' and color is not None:
 322             r = reset()
 323         else:
 324             r = ''
 325         return f'{left}[ {color}{title}{r} ]{right}'
 326     else:
 327         return ''
 328
 329
 330 def box(
 331     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
 332 ):
 333     """Draws a box with nice rounded corners.
 334
 335     >>> box('Title', 'This is text', width=30)
 336     ╭────────────────────────────╮
 337     │            Title           │
 338     │                            │
 339     │ This is text               │
 340     ╰────────────────────────────╯
 341
 342     """
 343     assert width > 4
 344     if color == '':
 345         rset = ''
 346     else:
 347         rset = reset()
 348     w = width - 2
 349     print(color + '╭' + '─' * w + '╮' + rset)
 350     if title is not None:
 351         print(
 352             color + '│' + rset + justify_string(title, width=w, alignment='c') + color + '│' + rset
 353         )
 354         print(color + '│' + ' ' * w + '│' + rset)
 355     if text is not None:
 356         for line in justify_text(text, width=w - 2, alignment='l').split('\n'):
 357             tw = len(line)
 358             assert tw < w
 359             print(color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset)
 360     print(color + '╰' + '─' * w + '╯' + rset)
 361
 362
 363 if __name__ == '__main__':
 364     import doctest
 365
 366     doctest.testmod()