path: root/text_utils.py

#!/usr/bin/env python3

"""Utilities for dealing with "text"."""

from collections import defaultdict
import logging
import math
import sys
from typing import List, NamedTuple, Optional

from ansi import fg, reset


logger = logging.getLogger(__file__)


class RowsColumns(NamedTuple):
    rows: int
    columns: int


def get_console_rows_columns() -> RowsColumns:
    """Returns the number of rows/columns on the current console."""

    from exec_utils import cmd_with_timeout
    try:
        rows, columns = cmd_with_timeout(
            "stty size",
            timeout_seconds=1.0,
        ).split()
    except Exception as e:
        logger.exception(e)
        raise Exception('Can\'t determine console size?!')
    return RowsColumns(int(rows), int(columns))


def progress_graph(
    current: int,
    total: int,
    *,
    width=70,
    fgcolor=fg("school bus yellow"),
    left_end="[",
    right_end="]",
    redraw=True,
) -> None:
    """Draws a progress graph."""

    percent = current / total
    ret = "\r" if redraw else "\n"
    bar = bar_graph(
        percent,
        include_text = True,
        width = width,
        fgcolor = fgcolor,
        left_end = left_end,
        right_end = right_end)
    print(
        bar,
        end=ret,
        flush=True,
        file=sys.stderr)


def bar_graph(
    percentage: float,
    *,
    include_text=True,
    width=70,
    fgcolor=fg("school bus yellow"),
    reset=reset(),
    left_end="[",
    right_end="]",
) -> None:
    """Returns a string containing a bar graph.

    >>> bar_graph(0.5, fgcolor='', reset='')
    '[███████████████████████████████████                                   ] 50.0%'

    """

    if percentage < 0.0 or percentage > 1.0:
        raise ValueError(percentage)
    if include_text:
        text = f"{percentage*100.0:2.1f}%"
    else:
        text = ""
    whole_width = math.floor(percentage * width)
    if whole_width == width:
        whole_width -= 1
        part_char = "▉"
    else:
        remainder_width = (percentage * width) % 1
        part_width = math.floor(remainder_width * 8)
        part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width]
    return (
        left_end +
        fgcolor +
        "█" * whole_width + part_char +
        " " * (width - whole_width - 1) +
        reset +
        right_end + " " +
        text)


def distribute_strings(
    strings: List[str],
    *,
    width: int = 80,
    alignment: str = "c",
    padding: str = " ",
) -> str:
    """
    Distributes strings into a line with a particular justification.

    >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
    '   this       is         a       test   '
    >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='l')
    'this      is        a         test      '
    >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='r')
    '      this        is         a      test'

    """
    subwidth = math.floor(width / len(strings))
    retval = ""
    for string in strings:
        string = justify_string(
            string, width=subwidth, alignment=alignment, padding=padding
        )
        retval += string
    while(len(retval) > width):
        retval = retval.replace('  ', ' ', 1)
    while(len(retval) < width):
        retval = retval.replace(' ', '  ', 1)
    return retval


def justify_string_by_chunk(
    string: str, width: int = 80, padding: str = " "
) -> str:
    """
    Justifies a string.

    >>> justify_string_by_chunk("This is a test", 40)
    'This       is              a        test'
    >>> justify_string_by_chunk("This is a test", 20)
    'This  is    a   test'

    """
    padding = padding[0]
    first, *rest, last = string.split()
    w = width - (len(first) + 1 + len(last) + 1)
    ret = first + padding + distribute_strings(rest, width=w, padding=padding) + padding + last
    return ret


def justify_string(
    string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
) -> str:
    """Justify a string.

    >>> justify_string('This is another test', width=40, alignment='c')
    '          This is another test          '
    >>> justify_string('This is another test', width=40, alignment='l')
    'This is another test                    '
    >>> justify_string('This is another test', width=40, alignment='r')
    '                    This is another test'
    >>> justify_string('This is another test', width=40, alignment='j')
    'This       is           another     test'

    """
    alignment = alignment[0]
    padding = padding[0]
    while len(string) < width:
        if alignment == "l":
            string += padding
        elif alignment == "r":
            string = padding + string
        elif alignment == "j":
            return justify_string_by_chunk(
                string,
                width=width,
                padding=padding
            )
        elif alignment == "c":
            if len(string) % 2 == 0:
                string += padding
            else:
                string = padding + string
        else:
            raise ValueError
    return string


def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
    """
    Justifies text.

    >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
    ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
    'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
    """
    retval = ""
    line = ""
    for word in text.split():
        if len(line) + len(word) > width:
            line = line[1:]
            line = justify_string(line, width=width, alignment=alignment)
            retval = retval + "\n" + line
            line = ""
        line = line + " " + word
    if len(line) > 0:
        retval += "\n" + line[1:]
    return retval[1:]


def generate_padded_columns(text: List[str]) -> str:
    max_width = defaultdict(int)
    for line in text:
        for pos, word in enumerate(line.split()):
            max_width[pos] = max(max_width[pos], len(word))

    for line in text:
        out = ""
        for pos, word in enumerate(line.split()):
            width = max_width[pos]
            word = justify_string(word, width=width, alignment='l')
            out += f'{word} '
        yield out


def wrap_string(text: str, n: int) -> str:
    chunks = text.split()
    out = ''
    width = 0
    for chunk in chunks:
        if width + len(chunk) > n:
            out += '\n'
            width = 0
        out += chunk + ' '
        width += len(chunk) + 1
    return out


class Indenter(object):
    """
    with Indenter(pad_count = 8) as i:
        i.print('test')
        with i:
            i.print('-ing')
            with i:
                i.print('1, 2, 3')
    """
    def __init__(self,
                 *,
                 pad_prefix: Optional[str] = None,
                 pad_char: str = ' ',
                 pad_count: int = 4):
        self.level = -1
        if pad_prefix is not None:
            self.pad_prefix = pad_prefix
        else:
            self.pad_prefix = ''
        self.padding = pad_char * pad_count

    def __enter__(self):
        self.level += 1
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.level -= 1
        if self.level < -1:
            self.level = -1

    def print(self, *arg, **kwargs):
        import string_utils
        text = string_utils.sprintf(*arg, **kwargs)
        print(self.pad_prefix + self.padding * self.level + text, end='')


def header(title: str, *, width: int = 80, color: str = ''):
    """
    Returns a nice header line with a title.

    >>> header('title', width=60, color='')
    '----[ title ]-----------------------------------------------'

    """
    w = width
    w -= (len(title) + 4)
    if w >= 4:
        left = 4 * '-'
        right = (w - 4) * '-'
        if color != '' and color is not None:
            r = reset()
        else:
            r = ''
        return f'{left}[ {color}{title}{r} ]{right}'
    else:
        return ''


if __name__ == '__main__':
    import doctest
    doctest.testmod()