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 dataclasses import dataclass
  12 from typing import Dict, Generator, List, Literal, Optional, Tuple
  13
  14 from ansi import fg, reset
  15
  16 logger = logging.getLogger(__file__)
  17
  18
  19 @dataclass
  20 class RowsColumns:
  21     """Row + Column"""
  22
  23     rows: int = 0
  24     columns: int = 0
  25
  26
  27 def get_console_rows_columns() -> RowsColumns:
  28     """Returns the number of rows/columns on the current console."""
  29
  30     from exec_utils import cmd
  31
  32     try:
  33         rows, columns = cmd(
  34             "stty size",
  35             timeout_seconds=1.0,
  36         ).split()
  37     except Exception as e:
  38         logger.exception(e)
  39         raise Exception('Can\'t determine console size?!') from e
  40     return RowsColumns(int(rows), int(columns))
  41
  42
  43 def progress_graph(
  44     current: int,
  45     total: int,
  46     *,
  47     width=70,
  48     fgcolor=fg("school bus yellow"),
  49     left_end="[",
  50     right_end="]",
  51     redraw=True,
  52 ) -> None:
  53     """Draws a progress graph."""
  54
  55     percent = current / total
  56     ret = "\r" if redraw else "\n"
  57     bar = bar_graph(
  58         percent,
  59         include_text=True,
  60         width=width,
  61         fgcolor=fgcolor,
  62         left_end=left_end,
  63         right_end=right_end,
  64     )
  65     print(bar, end=ret, flush=True, file=sys.stderr)
  66
  67
  68 def bar_graph(
  69     percentage: float,
  70     *,
  71     include_text=True,
  72     width=70,
  73     fgcolor=fg("school bus yellow"),
  74     reset=reset(),
  75     left_end="[",
  76     right_end="]",
  77 ) -> str:
  78     """Returns a string containing a bar graph.
  79
  80     >>> bar_graph(0.5, fgcolor='', reset='')
  81     '[███████████████████████████████████                                   ] 50.0%'
  82
  83     """
  84
  85     if percentage < 0.0 or percentage > 1.0:
  86         raise ValueError(percentage)
  87     if include_text:
  88         text = f"{percentage*100.0:2.1f}%"
  89     else:
  90         text = ""
  91     whole_width = math.floor(percentage * width)
  92     if whole_width == width:
  93         whole_width -= 1
  94         part_char = "▉"
  95     elif whole_width == 0 and percentage > 0.0:
  96         part_char = "▏"
  97     else:
  98         remainder_width = (percentage * width) % 1
  99         part_width = math.floor(remainder_width * 8)
 100         part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width]
 101     return (
 102         left_end
 103         + fgcolor
 104         + "█" * whole_width
 105         + part_char
 106         + " " * (width - whole_width - 1)
 107         + reset
 108         + right_end
 109         + " "
 110         + text
 111     )
 112
 113
 114 def sparkline(numbers: List[float]) -> Tuple[float, float, str]:
 115     """
 116     Makes a "sparkline" little inline histogram graph.  Auto scales.
 117
 118     >>> sparkline([1, 2, 3, 5, 10, 3, 5, 7])
 119     (1, 10, '▁▁▂▄█▂▄▆')
 120
 121     >>> sparkline([104, 99, 93, 96, 82, 77, 85, 73])
 122     (73, 104, '█▇▆▆▃▂▄▁')
 123
 124     """
 125     _bar = '▁▂▃▄▅▆▇█'  # Unicode: 9601, 9602, 9603, 9604, 9605, 9606, 9607, 9608
 126
 127     barcount = len(_bar)
 128     min_num, max_num = min(numbers), max(numbers)
 129     span = max_num - min_num
 130     sline = ''.join(
 131         _bar[min([barcount - 1, int((n - min_num) / span * barcount)])] for n in numbers
 132     )
 133     return min_num, max_num, sline
 134
 135
 136 def distribute_strings(
 137     strings: List[str],
 138     *,
 139     width: int = 80,
 140     alignment: str = "c",
 141     padding: str = " ",
 142 ) -> str:
 143     """
 144     Distributes strings into a line with a particular justification.
 145
 146     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
 147     '   this       is         a       test   '
 148     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='l')
 149     'this      is        a         test      '
 150     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='r')
 151     '      this        is         a      test'
 152
 153     """
 154     subwidth = math.floor(width / len(strings))
 155     retval = ""
 156     for string in strings:
 157         string = justify_string(string, width=subwidth, alignment=alignment, padding=padding)
 158         retval += string
 159     while len(retval) > width:
 160         retval = retval.replace('  ', ' ', 1)
 161     while len(retval) < width:
 162         retval = retval.replace(' ', '  ', 1)
 163     return retval
 164
 165
 166 def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:
 167     """
 168     Justifies a string.
 169
 170     >>> justify_string_by_chunk("This is a test", 40)
 171     'This       is              a        test'
 172     >>> justify_string_by_chunk("This is a test", 20)
 173     'This  is    a   test'
 174
 175     """
 176     padding = padding[0]
 177     first, *rest, last = string.split()
 178     w = width - (len(first) + 1 + len(last) + 1)
 179     ret = first + padding + distribute_strings(rest, width=w, padding=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 if __name__ == '__main__':
 331     import doctest
 332
 333     doctest.testmod()