text_utils.py

   1 #!/usr/bin/env python3
   2
   3 """Utilities for dealing with "text"."""
   4
   5 from collections import defaultdict
   6 import logging
   7 import math
   8 import sys
   9 from typing import List, NamedTuple, Optional
  10
  11 from ansi import fg, reset
  12
  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_with_timeout
  26     try:
  27         rows, columns = cmd_with_timeout(
  28             "stty size",
  29             timeout_seconds=1.0,
  30         ).split()
  31     except Exception as e:
  32         logger.exception(e)
  33         raise Exception('Can\'t determine console size?!')
  34     return RowsColumns(int(rows), int(columns))
  35
  36
  37 def progress_graph(
  38     current: int,
  39     total: int,
  40     *,
  41     width=70,
  42     fgcolor=fg("school bus yellow"),
  43     left_end="[",
  44     right_end="]",
  45     redraw=True,
  46 ) -> None:
  47     """Draws a progress graph."""
  48
  49     percent = current / total
  50     ret = "\r" if redraw else "\n"
  51     bar = bar_graph(
  52         percent,
  53         include_text = True,
  54         width = width,
  55         fgcolor = fgcolor,
  56         left_end = left_end,
  57         right_end = right_end)
  58     print(
  59         bar,
  60         end=ret,
  61         flush=True,
  62         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 ) -> None:
  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     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 + part_char +
 100         " " * (width - whole_width - 1) +
 101         reset +
 102         right_end + " " +
 103         text)
 104
 105
 106 def distribute_strings(
 107     strings: List[str],
 108     *,
 109     width: int = 80,
 110     alignment: str = "c",
 111     padding: str = " ",
 112 ) -> str:
 113     """
 114     Distributes strings into a line with a particular justification.
 115
 116     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
 117     '   this       is         a       test   '
 118     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='l')
 119     'this      is        a         test      '
 120     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='r')
 121     '      this        is         a      test'
 122
 123     """
 124     subwidth = math.floor(width / len(strings))
 125     retval = ""
 126     for string in strings:
 127         string = justify_string(
 128             string, width=subwidth, alignment=alignment, padding=padding
 129         )
 130         retval += string
 131     while(len(retval) > width):
 132         retval = retval.replace('  ', ' ', 1)
 133     while(len(retval) < width):
 134         retval = retval.replace(' ', '  ', 1)
 135     return retval
 136
 137
 138 def justify_string_by_chunk(
 139     string: str, width: int = 80, padding: str = " "
 140 ) -> str:
 141     """
 142     Justifies a string.
 143
 144     >>> justify_string_by_chunk("This is a test", 40)
 145     'This       is              a        test'
 146     >>> justify_string_by_chunk("This is a test", 20)
 147     'This  is    a   test'
 148
 149     """
 150     padding = padding[0]
 151     first, *rest, last = string.split()
 152     w = width - (len(first) + 1 + len(last) + 1)
 153     ret = first + padding + distribute_strings(rest, width=w, padding=padding) + padding + last
 154     return ret
 155
 156
 157 def justify_string(
 158     string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
 159 ) -> str:
 160     """Justify a string.
 161
 162     >>> justify_string('This is another test', width=40, alignment='c')
 163     '          This is another test          '
 164     >>> justify_string('This is another test', width=40, alignment='l')
 165     'This is another test                    '
 166     >>> justify_string('This is another test', width=40, alignment='r')
 167     '                    This is another test'
 168     >>> justify_string('This is another test', width=40, alignment='j')
 169     'This       is           another     test'
 170
 171     """
 172     alignment = alignment[0]
 173     padding = padding[0]
 174     while len(string) < width:
 175         if alignment == "l":
 176             string += padding
 177         elif alignment == "r":
 178             string = padding + string
 179         elif alignment == "j":
 180             return justify_string_by_chunk(
 181                 string,
 182                 width=width,
 183                 padding=padding
 184             )
 185         elif alignment == "c":
 186             if len(string) % 2 == 0:
 187                 string += padding
 188             else:
 189                 string = padding + string
 190         else:
 191             raise ValueError
 192     return string
 193
 194
 195 def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
 196     """
 197     Justifies text.
 198
 199     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
 200     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
 201     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
 202     """
 203     retval = ""
 204     line = ""
 205     for word in text.split():
 206         if len(line) + len(word) > width:
 207             line = line[1:]
 208             line = justify_string(line, width=width, alignment=alignment)
 209             retval = retval + "\n" + line
 210             line = ""
 211         line = line + " " + word
 212     if len(line) > 0:
 213         retval += "\n" + line[1:]
 214     return retval[1:]
 215
 216
 217 def generate_padded_columns(text: List[str]) -> str:
 218     max_width = defaultdict(int)
 219     for line in text:
 220         for pos, word in enumerate(line.split()):
 221             max_width[pos] = max(max_width[pos], len(word))
 222
 223     for line in text:
 224         out = ""
 225         for pos, word in enumerate(line.split()):
 226             width = max_width[pos]
 227             word = justify_string(word, width=width, alignment='l')
 228             out += f'{word} '
 229         yield out
 230
 231
 232 def wrap_string(text: str, n: int) -> str:
 233     chunks = text.split()
 234     out = ''
 235     width = 0
 236     for chunk in chunks:
 237         if width + len(chunk) > n:
 238             out += '\n'
 239             width = 0
 240         out += chunk + ' '
 241         width += len(chunk) + 1
 242     return out
 243
 244
 245 class Indenter(object):
 246     """
 247     with Indenter(pad_count = 8) as i:
 248         i.print('test')
 249         with i:
 250             i.print('-ing')
 251             with i:
 252                 i.print('1, 2, 3')
 253     """
 254     def __init__(self,
 255                  *,
 256                  pad_prefix: Optional[str] = None,
 257                  pad_char: str = ' ',
 258                  pad_count: int = 4):
 259         self.level = -1
 260         if pad_prefix is not None:
 261             self.pad_prefix = pad_prefix
 262         else:
 263             self.pad_prefix = ''
 264         self.padding = pad_char * pad_count
 265
 266     def __enter__(self):
 267         self.level += 1
 268         return self
 269
 270     def __exit__(self, exc_type, exc_value, exc_tb):
 271         self.level -= 1
 272         if self.level < -1:
 273             self.level = -1
 274
 275     def print(self, *arg, **kwargs):
 276         import string_utils
 277         text = string_utils.sprintf(*arg, **kwargs)
 278         print(self.pad_prefix + self.padding * self.level + text, end='')
 279
 280
 281 def header(title: str, *, width: int = 80, color: str = ''):
 282     """
 283     Returns a nice header line with a title.
 284
 285     >>> header('title', width=60, color='')
 286     '----[ title ]-----------------------------------------------'
 287
 288     """
 289     w = width
 290     w -= (len(title) + 4)
 291     if w >= 4:
 292         left = 4 * '-'
 293         right = (w - 4) * '-'
 294         if color != '' and color is not None:
 295             r = reset()
 296         else:
 297             r = ''
 298         return f'{left}[ {color}{title}{r} ]{right}'
 299     else:
 300         return ''
 301
 302
 303 if __name__ == '__main__':
 304     import doctest
 305     doctest.testmod()