text_utils.py

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