kglobe/kitty.py

#!/usr/bin/env python3
import asyncio
import re
import sys
import termios
import tty
from base64 import standard_b64encode
from io import BytesIO
import array
import fcntl


async def draw_to_terminal(buffer: BytesIO) -> None:
    """Display a PNG image in the terminal."""
    await _write_chunked(a="T", i=1, f=100, q=2, data=buffer.getvalue())


async def get_terminal_cell_size() -> tuple[int, int]:
    """Get (height, width) of a single cell in px"""
    reply = await _query_terminal("\x1b[16t", "t")

    match = re.search(r"\[6;(\d+);(\d+)t", reply)
    if match:
        v_pix, h_pix = map(int, match.groups())
        return v_pix, h_pix
    print(reply)
    raise ValueError("Failed to parse terminal cell size response")


async def get_terminal_size_pixel() -> tuple[int, int]:
    """Get (height, width) of the terminal in px"""
    reply = await _query_terminal("\x1b[14t", "t")

    match = re.search(r"\[4;(\d+);(\d+)t", reply)
    if match:
        height, width = map(int, match.groups())
        return height, width
    raise ValueError("Failed to parse terminal pixel size response")


def get_terminal_size() -> tuple[int, int]:
    """Get (rows, cols) of the terminal"""
    buf = array.array("H", [0, 0, 0, 0])
    fcntl.ioctl(sys.stdout, termios.TIOCGWINSZ, buf)
    rows, cols, width, height = buf
    return (rows, cols)


def _serialize_gr_command(**cmd) -> bytes:
    payload = cmd.pop("payload", None)
    cmd = ",".join(f"{k}={v}" for k, v in cmd.items())
    ans = []
    w = ans.append
    w(b"\033_G"), w(cmd.encode("ascii"))
    if payload:
        w(b";")
        w(payload)
    w(b"\033\\")
    return b"".join(ans)


async def _write_chunked(**cmd) -> None:
    data = standard_b64encode(cmd.pop("data"))
    loop = asyncio.get_running_loop()
    while data:
        chunk, data = data[:4096], data[4096:]
        m = 1 if data else 0
        payload = _serialize_gr_command(payload=chunk, m=m, **cmd)
        await loop.run_in_executor(None, sys.stdout.buffer.write, payload)
        await loop.run_in_executor(None, sys.stdout.flush)
        cmd.clear()


def hide_cursor() -> None:
    """Tell the terminal to hide the cursor."""
    _write_stdout("\x1b[?25l")


def show_cursor() -> None:
    """Tell the terminal to show the cursor."""
    _write_stdout("\x1b[?25h")


def set_position(y: int, x: int) -> None:
    """Set the cursor position to y, x"""
    _write_stdout(f"\x1b[{y};{x}H")


def _write_stdout(cmd: str) -> None:
    sys.stdout.write(cmd)
    sys.stdout.flush()


async def _query_terminal(escape: str, endchar: str) -> str:
    """
    Send `escape` to the terminal, read the response until
    `endchar`, return response (including `endchar`)
    """
    # Save the current terminal settings
    loop = asyncio.get_running_loop()
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)

    def read_terminal_response() -> str:
        tty.setraw(fd)
        _write_stdout(escape)
        response = ""
        while True:
            char = sys.stdin.read(1)
            response += char
            if char == endchar:
                break
        return response

    try:
        response = await loop.run_in_executor(None, read_terminal_response)
    finally:
        # Restore the terminal settings
        termios.tcsetattr(fd, termios.TCSANOW, old_settings)
    return response


async def get_position() -> tuple[int, int]:
    """Get the (y, x) position of the cursor"""
    reply = await _query_terminal("\x1b[6n", "R")

    match = re.search(r"\[(\d+);(\d+)R", reply)
    if match:
        y, x = map(int, match.groups())
        return y, x
    raise ValueError("Failed to parse cursor position response")


async def main():
    import pyvista as pv

    rows, cols = get_terminal_size()
    v_pix, h_pix = await get_terminal_cell_size()
    height, width = await get_terminal_size_pixel()
    y, x = await get_position()

    y, x = await get_position()
    print(f"Terminal has {rows} rows, {cols} cols = {rows * cols} cells")
    print(f"Cell size: {h_pix}x{v_pix} px")
    print(f"Dimensions: {width}x{height} px")
    print("Cursor is at y, x:", y, x)

    new_lines = int((height / v_pix) - 6)
    print("\n" * new_lines, end="")

    set_position(y - new_lines - 6, x)

    s = pv.Sphere()
    pl = pv.Plotter(off_screen=True)
    pl.add_mesh(s)
    b = BytesIO()
    pl.screenshot(b, transparent_background=True, window_size=(width, height))

    await draw_to_terminal(b)

if __name__ == "__main__":
    asyncio.run(main())