path: root/text_utils.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

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

import contextlib
import logging
import math
import sys
from collections import defaultdict
from typing import Dict, Generator, List, Literal, NamedTuple, Optional, Tuple

from ansi import fg, reset

logger = logging.getLogger(__file__)


class RowsColumns(NamedTuple):
    """Row + Column"""

    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

    try:
        rows, columns = cmd(
            "stty size",
            timeout_seconds=1.0,
        ).split()
    except Exception as e:
        logger.exception(e)
        raise Exception('Can\'t determine console size?!') from e
    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="]",
) -> str:
    """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 = "▉"
    elif whole_width == 0 and percentage > 0.0:
        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 sparkline(numbers: List[float]) -> Tuple[float, float, str]:
    """
    Makes a "sparkline" little inline histogram graph.  Auto scales.

    >>> sparkline([1, 2, 3, 5, 10, 3, 5, 7])
    (1, 10, '▁▁▂▄█▂▄▆')

    >>> sparkline([104, 99, 93, 96, 82, 77, 85, 73])
    (73, 104, '█▇▆▆▃▂▄▁')

    """
    _bar = '▁▂▃▄▅▆▇█'  # Unicode: 9601, 9602, 9603, 9604, 9605, 9606, 9607, 9608

    barcount = len(_bar)
    min_num, max_num = min(numbers), max(numbers)
    span = max_num - min_num
    sline = ''.join(
        _bar[min([barcount - 1, int((n - min_num) / span * barcount)])] for n in numbers
    )
    return min_num, max_num, sline


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]) -> Generator:
    max_width: Dict[int, int] = 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(contextlib.AbstractContextManager):
    """
    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) -> Literal[False]:
        self.level -= 1
        if self.level < -1:
            self.level = -1
        return False

    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()