tmenu.py

#!/home/mukundan/Documents/misc/tmenu/.venv/bin/python

import base64
import curses
import curses.ascii
import fcntl
import itertools
import os
import struct
import subprocess
import sys
import termios

import xdg.BaseDirectory
import xdg.IconTheme
import xdg.Menu

image_id_generator = itertools.count(1)
default_icon_path = "/usr/share/icons/Papirus-Dark/symbolic/mimetypes/application-x-executable-symbolic.svg"
menu_icon_path = "/usr/share/icons/Papirus-Dark/symbolic/actions/open-menu-symbolic.svg"


def truncate(text: str, size: int):
    return text if len(text) <= size else text[:size - 3] + "..."


def get_cell_size() -> tuple[int, int]:
    ws_row, ws_col, ws_xpixel, ws_ypixel = struct.unpack(
        "HHHH",
        fcntl.ioctl(
            sys.stdin.fileno(),
            termios.TIOCGWINSZ,
            struct.pack("HHHH", 0, 0, 0, 0),
        )
    )

    cell_width = ws_xpixel // ws_col
    cell_height = ws_ypixel // ws_row

    return (cell_width, cell_height)


def load_image(path: str, size: tuple[int, int]) -> int:
    cell_size = get_cell_size()
    cache_path = os.path.join(
        xdg.BaseDirectory.save_cache_path("tmenu"),
        f"{path.replace('%', '%%').replace('/', '%')[-20:]}-{cell_size[0]}x{cell_size[1]}-{size[0]}x{size[1]}.png"
    )

    if not os.path.isfile(cache_path):
        subprocess.run(
            [
                "magick",
                "convert",
                "-background",
                "none",
                "-density",
                "500",
                "-resize",
                f"{int(size[1]*cell_size[0]*0.8)}x{(size[0]*cell_size[1]*0.8)}",
                path,
                "-gravity",
                "center",
                "-extent",
                f"{size[1]*cell_size[0]}x{size[0]*cell_size[1]}",
                cache_path,
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )

    i = next(image_id_generator)
    print(end=f"\033_Ga=t,f=100,t=f,q=1,i={i};{base64.b64encode(cache_path.encode()).decode()}\033\\", flush=True)

    return i


def display_image(win: curses.window, y: int, x: int, i: int):
    subwin = win.derwin(1, 1, y, x)
    subwin.move(0, 0)
    subwin.refresh()
    print(end=f"\033_Ga=d,d=c;\033\\\033_Ga=p,i={i},q=1,C=1;\033\\", flush=True)


def delete_image(win: curses.window, y: int, x: int):
    subwin = win.derwin(1, 1, y, x)
    subwin.move(0, 0)
    subwin.refresh()
    print(end="\033_Ga=d,d=c;\033\\", flush=True)


def delete_images(free=False):
    print(f"\033_Ga=d{',d=A' if free else ''};\033\\", end="", flush=True)


def load_menu_icons(menu):
    size = min(*get_cell_size())
    menu_icons = dict()
    default_icon = load_image(default_icon_path, (1, 4))
    menu_icon = load_image(menu_icon_path, (1, 4))
    for entry in menu.getEntries():
        if isinstance(entry, xdg.Menu.MenuEntry):
            icon = entry.DesktopEntry.getIcon()
            path = xdg.IconTheme.getIconPath(entry.DesktopEntry.getIcon(), size=size, theme='Papirus-Dark') # TODO: cache this
            if path and icon not in menu_icons:
                menu_icons[icon] = load_image(path, (1, 4))

    def get_icon(entry):
        if isinstance(entry, xdg.Menu.Menu):
            return menu_icon
        elif isinstance(entry, xdg.Menu.MenuEntry):
            return menu_icons.get(entry.DesktopEntry.getIcon(), default_icon)

    return get_icon


def get_entry_name(entry):
    if isinstance(entry, xdg.Menu.Menu):
        return entry.getName()
    elif isinstance(entry, xdg.Menu.MenuEntry):
        return entry.DesktopEntry.getName()
    raise TypeError(f"must be xdg.Menu.Menu or xdg.Menu.MenuEntry, not {type(entry).__name__}")


def match(entry, query):
    return query.lower() in get_entry_name(entry).lower()


def xdg_menu_choose(win, menu, prompt):
    size = win.getmaxyx()

    win.addstr(0, 0, prompt)

    get_icon = load_menu_icons(menu)

    # prompt_win = win.derwin(1, len(prompt) + 1, 0, 0)
    query_win = win.derwin(1, 0, 0, len(prompt))
    scrollbar_win = win.derwin(0, 0, 1, size[1] - 1)
    entry_text_win = win.derwin(0, size[1] - 5, 1, 4)
    entry_icon_win = win.derwin(0, 4, 1, 0)

    page_start = 0
    selected_idx = 0
    selected = None
    query = ""

    while True:
        entries = [entry for entry in menu.getEntries() if match(entry, query)]
        max_entry_text_size = entry_text_win.getmaxyx()[1]
        selected_idx = max(0, min(len(entries) - 1, selected_idx))
        selected = entries[selected_idx] if entries else None
        page_size = entry_text_win.getmaxyx()[0]
        page_start = max(0, selected_idx - page_size + 1, min(len(entries) - page_size, page_start, selected_idx))
        page_end = min(len(entries), page_start + page_size)

        query_win.erase()
        query_win.addstr(0, 0, query[-query_win.getmaxyx()[1] + 1:])
        query_win.noutrefresh()

        entry_text_win.erase()
        entry_icon_win.erase()
        for row, entry in enumerate(entries[page_start:page_end]):
            idx = row + page_start
            entry_text_win.addstr(row, 0, truncate(get_entry_name(entry), max_entry_text_size))
            if idx == selected_idx:
                entry_text_win.chgat(row, 0, curses.A_STANDOUT)
                entry_icon_win.chgat(row, 0, curses.A_STANDOUT)
        entry_text_win.noutrefresh()
        entry_icon_win.noutrefresh()

        scrollbar_win.erase()
        scrollbar_size = scrollbar_win.getmaxyx()[0] * 8
        scrollthumb_start = round(scrollbar_size * (page_start / len(entries))) if entries else 0
        scrollthumb_end = round(scrollbar_size * (page_end / len(entries))) if entries else scrollbar_size
        for row in range(scrollbar_win.getmaxyx()[0]):
            if row * 8 <= scrollthumb_start <= (row + 1) * 8:
                scrollbar_win.insstr(row, 0, "█▇▆▅▄▃▂▁ "[scrollthumb_start - (row * 8)])
            if row * 8 <= scrollthumb_end <= (row + 1) * 8:
                scrollbar_win.insstr(row, 0, " ▔🮂🮃▀🮄🮅🮆█"[scrollthumb_end - (row * 8)])
            if scrollthumb_start <= row * 8 and (row + 1) * 8 <= scrollthumb_end:
                scrollbar_win.insstr(row, 0, "█")
        scrollbar_win.noutrefresh()

        win.refresh()

        # display images after other changes to avoid curses from deleting the image accidentally
        curses.curs_set(0)
        for row in range(page_size):
            idx = row + page_start
            if row < len(entries):
                display_image(entry_icon_win, row, 0, get_icon(entries[idx]))
            else:
                delete_image(entry_icon_win, row, 0)
        curses.curs_set(1)

        query_win.cursyncup()
        match win.getch():
            case curses.KEY_DOWN | curses.KEY_SF:
                selected_idx += 1
            case curses.KEY_UP | curses.KEY_SR:
                selected_idx -= 1
            case curses.KEY_RIGHT | curses.KEY_ENTER | curses.ascii.LF:
                break
            case curses.KEY_LEFT | curses.ascii.ESC:
                selected = None
                break
            case curses.KEY_RESIZE:
                size = win.getmaxyx()
                # prompt_win = win.derwin(1, len(prompt) + 1, 0, 0)
                query_win = win.derwin(1, 0, 0, len(prompt))
                scrollbar_win = win.derwin(0, 0, 1, size[1] - 1)
                entry_text_win = win.derwin(0, size[1] - 5, 1, 4)
                entry_icon_win = win.derwin(0, 4, 1, 0)
                delete_images(True)
                get_icon = load_menu_icons(menu)
            case curses.KEY_BACKSPACE:
                query = query[:-1]
            case ch if curses.ascii.isprint(ch):
                query += chr(ch)

    # feels more responsive if everything is removed rather than just images
    win.erase()
    win.refresh()
    delete_images(True)

    return selected


def main(stdscr: curses.window):
    curses.set_escdelay(10)

    prev_menus = []
    current_menu = xdg.Menu.parse()
    while True:
        choice = xdg_menu_choose(stdscr, current_menu, f" {current_menu.getName()} > ")
        if isinstance(choice, xdg.Menu.Menu):
            prev_menus.append(current_menu)
            current_menu = choice
        elif isinstance(choice, xdg.Menu.MenuEntry):
            subprocess.run([
                "swaymsg",
                "exec",
                "gio",
                "launch",
                choice.DesktopEntry.filename,
            ])
            break
        elif choice is None:
            if not prev_menus:
                break
            current_menu = prev_menus.pop()


if __name__ == "__main__":
    curses.wrapper(main)