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