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