pt_live.py

#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "ptpython",
#     "rich",
# ]
# ///


import contextlib
import logging
import os
import pkgutil
import select
import sys
import termios
import tty
from pathlib import Path
from subprocess import run

from ptpython.entry_points import run_ptpython  # type: ignore
from rich.live import Live
from rich.text import Text

logger = logging.getLogger(__name__)
logger.setLevel(logging.WARNING)


class LiveMenu:
    """Autocompletion menu."""

    def __init__(
        self,
        cmd_line: str,
        modules: list[str],
        prefix: str | None = None,
        cmd_point: int | None = None,
        menu_height: int = 5,
    ) -> None:
        self._fd = sys.stdin.fileno()
        if not os.isatty(self._fd):
            raise OSError("Input not from tty.")

        try:
            self._menu_height = min(
                menu_height, int(os.get_terminal_size(self._fd)[1]) - 1
            )
        except (OSError, ValueError):
            self._menu_height = menu_height

        self.modules = modules
        self._selected = -1
        self._items = []
        self._start_idx = 0
        self._stop_idx = self._start_idx + self._menu_height
        self._sep = [" ", ","]
        self._prefix = (
            Text.from_ansi("") if prefix is None else Text.from_ansi(prefix)
        )
        self._cmd_line = cmd_line
        self._cmd_point = cmd_point or len(cmd_line)
        self._cmd_style = "color(247)"
        self._word = ""
        self._completed = True
        self._menu_style = "color(66) on color(235)"
        self._menu_sel_style = "bold magenta on color(237)"

    def __enter__(self):
        self._old = termios.tcgetattr(self._fd)
        tty.setcbreak(self._fd)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old)
        return False

    def _generate_items(self, line: str | None, idx: int) -> list[str]:
        """Generate list of completion items."""
        if idx > 0 and (line is None or idx > len(line)):
            raise IndexError(f"{idx=} is invalid for {line=}")

        word = ""
        module_list = self.modules[:]
        if line:
            if idx:
                i = 1
                while i <= idx and line[idx - i].isalpha():
                    word = line[idx - i] + word
                    i += 1

            for module in self.modules:
                if module in line.split():
                    with contextlib.suppress(ValueError):
                        module_list.remove(module)

        if word:
            for module in self.modules:
                if not module.startswith(word):
                    with contextlib.suppress(ValueError):
                        module_list.remove(module)

        return module_list

    def get_key(self) -> str | None:
        termios.tcflush(self._fd, termios.TCIFLUSH)
        rlist = select.select([self._fd], [], [], 50)[0]
        if rlist:
            return os.read(self._fd, 6).decode()

        return None

    def prepare_menu(self) -> Text:
        logger.debug(
            f"prepare_menu:<{self._cmd_line=}> : <{self._cmd_point=}>"
        )
        self._items = self._generate_items(self._cmd_line, self._cmd_point)

        text = self._prefix.copy()
        indent = " " * (text.cell_len + self._cmd_point - len(self._word))
        length = (
            max(
                len(line)
                for line in self._items[self._start_idx : self._stop_idx]
            )
            if self._items
            else 0
        )
        head = self._cmd_line[: self._cmd_point]
        tail = self._cmd_line[self._cmd_point :] or "┊"
        if self._selected != -1:
            head = head.removesuffix(self._word) + self._items[self._selected]

        text.append(
            Text.from_markup(
                f"[{self._cmd_style}]{head}[reverse]{tail[:1]}[/]{tail[1:]}[/]"
            )
        )

        if self._completed:
            self._completed = False
        else:
            for i, item in enumerate(
                self._items[self._start_idx : self._stop_idx],
                start=self._start_idx,
            ):
                if i == self._selected:
                    style = self._menu_sel_style
                else:
                    style = self._menu_style

                text.append(
                    Text.from_markup(
                        f"\n{indent}[{style}]{item:<{length}}[/]"
                    )
                )

        return text

    def display_menu(self) -> str:
        with Live(
            self.prepare_menu(), transient=True, auto_refresh=False
        ) as live:
            while True:
                result = self._parse_input()
                if result is not None:
                    return result

                live.update(self.prepare_menu(), refresh=True)

    def _parse_input(self) -> str | None:
        key = self.get_key()

        # TAB, Down Arrow, Ctrl-n
        if key in ["\t", "\x1b[B", "\x0e"]:
            logger.debug("TAB/Down pressed")
            if self._selected < len(self._items) - 1:
                self._selected += 1
                if self._selected > self._stop_idx - 1:
                    self._stop_idx = self._selected + 1
                    self._start_idx = self._stop_idx - self._menu_height
            else:
                self._selected = -1
                self._start_idx = 0
                self._stop_idx = self._menu_height

        # Shift-TAB or Up Arrow
        elif key in ["\x1b[Z", "\x1b[A", "\x10"]:
            logger.debug("Shift-TAB/Up pressed")
            if self._selected > -1:
                self._selected -= 1
                if self._selected < self._start_idx:
                    self._start_idx = max(self._selected, 0)
                    self._stop_idx = self._start_idx + self._menu_height
            else:
                self._selected = len(self._items) - 1
                self._stop_idx = self._selected + 1
                self._start_idx = self._stop_idx - self._menu_height

        # Escape
        elif key == "\x1b":
            sys.exit()

        # Enter, Alt-p, Ctrl-m, Ctrl-d
        elif key in ["\n", "\x1bp", "\x0d", "\x04"]:
            logger.debug("Enter/Escape pressed")
            if self._selected != -1:
                head = self._cmd_line[: self._cmd_point]
                tail = self._cmd_line[self._cmd_point :]
                if self._word and self._items[self._selected].startswith(
                    self._word
                ):
                    head = head.removesuffix(self._word)
                    self._cmd_point -= len(self._word)
                head += f"{self._items[self._selected]}"
                self._cmd_line = head + tail
                self._cmd_point += len(self._items[self._selected])
                self._word = ""
                self._completed = True
                self._selected = -1
                self._start_idx = 0
                self._stop_idx = self._start_idx + self._menu_height
            else:
                return self._cmd_line

        # elif " " <= key <= "~":
        elif key and key.isprintable():
            logger.debug("key printable")

            head = self._cmd_line[: self._cmd_point]
            tail = self._cmd_line[self._cmd_point :]
            if key not in self._sep:
                self._word += key
            elif self._selected != -1:
                head = (
                    head.removesuffix(self._word)
                    + self._items[self._selected]
                )
                self._cmd_point += len(self._items[self._selected]) - len(
                    self._word
                )
                self._word = ""
                self._selected = -1
                self._start_idx = 0
                self._stop_idx = self._start_idx + self._menu_height
            else:
                self._word = ""

            self._cmd_line = head + key + tail
            self._cmd_point += 1

        # Backspace
        elif key == "\x7f":
            logger.debug("Backspace  pressed")
            head = self._cmd_line[: self._cmd_point]
            tail = self._cmd_line[self._cmd_point :]
            if self._selected != -1:
                # Replace word with selected item, than backtrack
                head = (
                    head.removesuffix(self._word)
                    + self._items[self._selected]
                )
                self._cmd_point += len(self._items[self._selected]) - len(
                    self._word
                )
                self._word = self._items[self._selected]
                self._selected = -1
                self._start_idx = 0
                self._stop_idx = self._start_idx + self._menu_height

            if self._word:
                # Remove last character in word
                self._word = self._word[:-1]

            if self._cmd_point > 0:
                self._cmd_point -= 1

            # Remove character before point
            self._cmd_line = head[: self._cmd_point] + tail
            logger.debug(f"{self._word=}")

        else:
            return


class PtHistory:
    """Load/save list of lines from/to history"""

    def __init__(self, history_file_str: str) -> None:
        self.history_file = Path(history_file_str).expanduser()
        if self.history_file.exists():
            self.history_list = [
                line
                for line in self.history_file.read_text().splitlines()
                if line
            ]
        else:
            self.history_list = []

    def get_history(self) -> list[str]:
        """Return a copy of history list"""
        return self.history_list[:]

    def save_history(self, new_entry: str) -> None:
        """Save history list to file"""
        if new_entry:
            new_history = [new_entry]
            for line in self.history_list:
                if line not in new_history:
                    new_history.append(line)

            with self.history_file.open("w", encoding="utf-8") as f:
                f.write("\n".join(new_history))


def venv_get(search="python", start=".", fallback=None):
    if search.startswith("py"):
        glob = "*/bin/python"
    elif search.startswith("si"):
        glob = "*/lib/python*/site-packages"
    else:
        return None

    p = Path(start).expanduser().absolute()

    while str(p) != p.root:
        for path in p.glob(glob):
            return str(path)
        p = p.parent

    return fallback


def main():
    """Start ptpython repl."""

    logger.debug("__debug_started__")
    history_file = "~/.pti_history"
    history = PtHistory(history_file)
    modules = history.get_history()

    # Find virtual environment
    fallback = "~/code"
    venv_py = venv_get("python", fallback=fallback)
    site_pkgs = venv_get("site", fallback=fallback)

    # Get a list of modules from venv
    if site_pkgs:
        sys.path.append(site_pkgs)
        module_list = pkgutil.iter_modules([site_pkgs])
        for module in module_list:
            if not module.name.startswith("_") and module.name not in modules:
                modules.append(module.name)

    prefix = sys.argv[1] if len(sys.argv) == 2 else " => "

    cmd_line = ""

    # Choose modules
    with LiveMenu(
        cmd_line,
        modules,
        cmd_point=0,
        prefix=prefix,
    ) as menu:
        new_line = menu.display_menu()

    new_history_entry = new_line.strip().replace(" ", ",")
    if logger.isEnabledFor(logging.DEBUG):
        logger.debug("command = '%s'", new_line)
        logger.debug("__debug_ended__")
        return

    if new_history_entry:
        # Add modules to history file
        history.save_history(new_history_entry)
        cmd = ["uv", "pip", "install"]
        if venv_py:
            cmd.extend(["--python", venv_py])
        cmd.extend(new_history_entry.split(","))
        run(cmd, check=False)

    if len(sys.argv) > 1:
        sys.argv[:] = sys.argv[:1]

    run_ptpython.run()


if __name__ == "__main__":
    logging.basicConfig(
        filename="pt.log",
        format="%(asctime)s:%(levelname)s:%(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    main()