dnd/deadsands/site_tools/console.py

import os

from configparser import ConfigParser
from pathlib import Path
from textwrap import dedent
from typing import Union, List

import rich.repr

from rich.console import Console as _Console
from rich.markdown import Markdown
from rich.theme import Theme
from rich.table import Table, Column

from prompt_toolkit import PromptSession
from prompt_toolkit.patch_stdout import patch_stdout
from prompt_toolkit.formatted_text import ANSI
from prompt_toolkit.styles import Style
from prompt_toolkit.output import ColorDepth


BASE_STYLE = {
    'help': 'cyan',
    'bright': 'white',
    'repr.str': 'dim',
    'repr.brace': 'dim',
    'repr.url': 'blue',
    'table.header': 'white',
    'toolbar.fg': '#888888',
    'toolbar.bg': '#111111',
    'toolbar.bold': '#FFFFFF',
    'error': 'red',
}


def console_theme(theme_name: Union[str, None] = None) -> dict:
    """
    Return a console theme as a dictionary.

    Args:
        theme_name (str):
    """
    cfg = ConfigParser()
    cfg.read_dict({'styles': BASE_STYLE})

    if theme_name:
        theme_path = theme_name if theme_name else os.environ.get('DEFAULT_THEME', 'blue_train')
        cfg.read(Theme(
            Path(theme_path) / Path('console.cfg')
        ))
    return cfg['styles']


@rich.repr.auto
class Console(_Console):
    """
    SYNOPSIS

        Subclasses a rich.console.Console to provide an instance with a
        reconfigured themes, and convenience methods and attributes.

    USAGE

        Console([ARGS])

    ARGS

        theme       The name of a theme to load. Defaults to DEFAULT_THEME.

    EXAMPLES

        Console().print("Can I kick it?")
        >>> Can I kick it?

    INSTANCE ATTRIBUTES

        theme       The current theme

    """

    def __init__(self, *args, **kwargs):
        self._console_theme = console_theme(kwargs.get('theme', None))
        self._overflow = 'ellipsis'
        kwargs['theme'] = Theme(self._console_theme, inherit=False)
        super().__init__(*args, **kwargs)

        self._session = PromptSession()

    @property
    def theme(self) -> Theme:
        return self._console_theme

    def prompt(self, lines: List, **kwargs) -> str:
        """
        Print a list of lines, using the final line as a prompt.

        Example:

            Console().prompt(["Can I kick it?", "[Y/n] ")
            >>> Can I kick it?
            [Y/n]>

        """

        prompt_style = Style.from_dict({
            # 'bottom-toolbar': f"{self.theme['toolbar.fg']} bg:{self.theme['toolbar.bg']}",
            # 'toolbar-bold': f"{self.theme['toolbar.bold']}"
        })

        for line in lines[:-1]:
            super().print(line)
        with self.capture() as capture:
            super().print(f"[prompt bold]{lines[-1]}>[/] ", end='')
        text = ANSI(capture.get())

        # This is required to intercept key bindings and not mess up the
        # prompt. Both the prompt and bottom_toolbar params must be functions
        # for this to correctly regenerate the prompt after the interrupt.
        with patch_stdout(raw=True):
            return self._session.prompt(
                lambda: text,
                style=prompt_style,
                color_depth=ColorDepth.TRUE_COLOR,
                **kwargs)

    def mdprint(self, txt: str, **kwargs) -> None:
        """
        Like print(), but support markdown. Text will be dedented.
        """
        self.print(Markdown(dedent(txt), justify='left'), **kwargs)

    def print(self, txt: str, **kwargs) -> None:
        """
        Print text to the console, possibly truncated with an ellipsis.
        """
        super().print(txt, overflow=self._overflow, **kwargs)

    def debug(self, txt: str, **kwargs) -> None:
        """
        Print text to the console with the current theme's debug style applied, if debugging is enabled.
        """
        if os.environ.get('DEBUG', None):
            self.print(dedent(txt), style='debug')

    def error(self, txt: str, **kwargs) -> None:
        """
        Print text to the console with the current theme's error style applied.
        """
        self.print(dedent(txt), style='error')

    def table(self, *cols: List[Column], **params) -> None:
        """
        Print a rich table to the console with theme elements and styles applied.
        parameters and keyword arguments are passed to rich.table.Table.
        """
        background_style = f"on {self.theme['background']}"
        params.update(
            header_style=background_style,
            title_style=background_style,
            border_style=background_style,
            row_styles=[background_style],
            caption_style=background_style,
            style=background_style,
        )
        params['min_width'] = 80
        width = os.environ.get('CONSOLE_WIDTH', 'auto')
        if width == 'expand':
            params['expand'] = True
        elif width != 'auto':
            params['width'] = int(width)
        return Table(*cols, **params)