diff --git a/kitty.py b/kitty.py index 36725d4..7cee3bb 100644 --- a/kitty.py +++ b/kitty.py @@ -1,21 +1,18 @@ #!/usr/bin/env python3 +from base64 import standard_b64encode +from enum import Enum +import array import asyncio +import fcntl +import mmap +import os import re +import select import sys import termios import tty -from base64 import standard_b64encode -import array -import fcntl - import zlib - -import mmap -import os -import select -from enum import Enum - from terminalio import TerminalIO SHM_NAME = "/kitty-shm-frame" @@ -26,6 +23,10 @@ class Kitty_Format(Enum): RGBA = 32 PNG = 100 +class TerminalProtocolError(Exception): + pass + + class Terminal(): def __init__(self): self.terminal_io = TerminalIO() @@ -42,7 +43,6 @@ class Terminal(): async def get_terminal_cell_size(self) -> tuple[int, int]: """Get (height, width) of a single cell in px""" reply = await self._query_terminal("\x1b[16t", "t") - match = re.search(r"\[6;(\d+);(\d+)t", reply) if match: v_pix, h_pix = map(int, match.groups()) @@ -54,7 +54,6 @@ class Terminal(): async def get_terminal_size_pixel(self) -> tuple[int, int]: """Get (height, width) of the terminal in px""" reply = await self._query_terminal("\x1b[14t", "t") - match = re.search(r"\[4;(\d+);(\d+)t", reply) if match: height, width = map(int, match.groups()) @@ -120,7 +119,7 @@ class Terminal(): # set shm name as payload data = SHM_NAME - await _write_chunked(**kwargs, data=data) + await _write_chunked(**kwargs, data=data, wait_for_ok=True if use_shm else False) await asyncio.sleep(0.001) async def _query_terminal(self, escape: str, endchar: str) -> str: @@ -128,7 +127,6 @@ class Terminal(): 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) @@ -151,6 +149,24 @@ class Terminal(): termios.tcsetattr(fd, termios.TCSANOW, old_settings) return response + async def _write_chunked(self, wait_for_ok: bool = True, **cmd) -> None: + image_id = cmd['i'] + if cmd["t"] == "s": + data = standard_b64encode(bytes(cmd.pop("data"), "utf-8")) + else: + data = standard_b64encode(cmd.pop("data")) + while data: + chunk, data = data[:4096], data[4096:] + m = 1 if data else 0 + payload = _serialize_gr_command(payload=chunk, m=m, **cmd) + sys.stdout.buffer.write(payload) + sys.stdout.flush() + cmd.clear() + + if wait_for_ok: + # Wait for terminal to confirm OK response + await self.terminal_io.wait_for_ok() + def _mmap(np_image, width: int, height: int) -> None: """Write image data to shared memory""" shm_size = width * height * 4 # RGBA @@ -205,18 +221,59 @@ def _serialize_gr_command(**cmd) -> bytes: return b"".join(ans) -async def _write_chunked(**cmd) -> None: +#async def _write_chunked(**cmd) -> None: +# if cmd["t"] == "s": +# data = standard_b64encode(bytes(cmd.pop("data"), "utf-8")) +# else: +# data = standard_b64encode(cmd.pop("data")) +# while data: +# chunk, data = data[:4096], data[4096:] +# m = 1 if data else 0 +# payload = _serialize_gr_command(payload=chunk, m=m, **cmd) +# sys.stdout.buffer.write(payload) +# sys.stdout.flush() +# cmd.clear() + +async def _write_chunked(wait_for_ok: bool = True, **cmd) -> None: + image_id = cmd['i'] if cmd["t"] == "s": data = standard_b64encode(bytes(cmd.pop("data"), "utf-8")) else: data = standard_b64encode(cmd.pop("data")) - while data: - chunk, data = data[:4096], data[4096:] - m = 1 if data else 0 - payload = _serialize_gr_command(payload=chunk, m=m, **cmd) - sys.stdout.buffer.write(payload) - sys.stdout.flush() - cmd.clear() + + fd = sys.stdin.fileno() + old_settings = termios.tcgetattr(fd) + tty.setraw(fd) + + loop = asyncio.get_running_loop() + + try: + while data: + chunk, data = data[:4096], data[4096:] + m = 1 if data else 0 + payload = _serialize_gr_command(payload=chunk, m=m, **cmd) + sys.stdout.buffer.write(payload) + sys.stdout.flush() + cmd.clear() + + if wait_for_ok: + # Wait for terminal to confirm OK response + response = os.read(fd, 1024).decode("utf-8", "ignore") + if not response.startswith(f"\033_Gi={image_id};OK"): + return + #raise RuntimeError(f"Unexpected response: {repr(response)}") + #return response.startswith("\033_GOK") + # def _read_ok(): + # response = "" + # while not response.endswith("\033\\"): + # response += sys.stdin.read(1) + # return response + + # response = await loop.run_in_executor(None, _read_ok) + # if not response.startswith(f"\033_Gi={image_id};OK"): + # raise RuntimeError(f"Unexpected response: {repr(response)}") + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) def _write_stdout(cmd: str) -> None: """Write a command string to stdout and flush.""" diff --git a/terminalio.py b/terminalio.py index 995815d..4cc9463 100644 --- a/terminalio.py +++ b/terminalio.py @@ -56,7 +56,7 @@ class TerminalIO: self._buffer += char if self._buffer.endswith("\033\\"): #if "\033_GOK" in self._buffer: - if re.search('\033_G(i=\d+;)?OK', self._buffer): + if re.search('\033_G(i=\\d+;)?OK', self._buffer): self.ok_event.set() self._buffer = "" continue @@ -65,4 +65,5 @@ class TerminalIO: if not char.startswith("\033"): for callback in self.input_callbacks: callback(char) + await asyncio.sleep(0.01) diff --git a/terminalplotter.py b/terminalplotter.py index 625775b..331e191 100644 --- a/terminalplotter.py +++ b/terminalplotter.py @@ -29,12 +29,13 @@ class TerminalPlotter(pv.Plotter): self.set_background([0.0, 0.0, 0.0]) self.ren_win.SetAlphaBitPlanes(1) self.ren_win.SetMultiSamples(0) + self.terminal = kitty.Terminal() async def initialize(self): - h_pix, _ = await kitty.get_terminal_cell_size() - self.rows, _ = kitty.get_terminal_size() - self.start_y, self.start_x = await kitty.get_position() + h_pix, _ = await self.terminal.get_terminal_cell_size() + self.rows, _ = self.terminal.get_terminal_size() + self.start_y, self.start_x = await self.terminal.get_position() self.needed_lines = math.ceil(self.height / h_pix) # Add vertical space if needed @@ -42,41 +43,61 @@ class TerminalPlotter(pv.Plotter): missing = self.needed_lines - (self.rows - self.start_y) self.start_y -= missing print("\n" * self.needed_lines, end="") - kitty.set_position(self.start_y, self.start_x) + self.terminal.set_position(self.start_y, self.start_x) self.orbit = self.generate_orbital_path() - async def _handle_key(self, c): - if c == "a": - self.camera.Azimuth(10) - elif c == "d": - self.camera.Azimuth(-10) - elif c == "w": - self.camera.Elevation(10) - elif c == "s": - self.camera.Elevation(-10) - elif c == "+": - self.camera.zoom(1.1) - elif c == "-": - self.camera.zoom(0.9) - elif c == "p": - self.fly_to(self.orbit.points[15]) + def _handle_key(c): + if c == "a": + self.camera.Azimuth(10) + elif c == "d": + self.camera.Azimuth(-10) + elif c == "w": + self.camera.Elevation(10) + elif c == "s": + self.camera.Elevation(-10) + elif c == "+": + self.camera.zoom(1.1) + elif c == "-": + self.camera.zoom(0.9) + elif c == "p": + self.fly_to(self.orbit.points[15]) + elif c == "q": + self._running = False + self.terminal.terminal_io.register_input_callback(_handle_key) - async def _input_loop(self): - fd = sys.stdin.fileno() - old_settings = termios.tcgetattr(fd) - try: - tty.setcbreak(fd) - while self._running: - rlist, _, _ = select.select([sys.stdin], [], [], 0) - if rlist: - c = sys.stdin.read(1) - if c == "q": # quit on q - self._running = False - else: - await self._handle_key(c) - await asyncio.sleep(0.001) - finally: - termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) + +# async def _handle_key(self, c): +# if c == "a": +# self.camera.Azimuth(10) +# elif c == "d": +# self.camera.Azimuth(-10) +# elif c == "w": +# self.camera.Elevation(10) +# elif c == "s": +# self.camera.Elevation(-10) +# elif c == "+": +# self.camera.zoom(1.1) +# elif c == "-": +# self.camera.zoom(0.9) +# elif c == "p": +# self.fly_to(self.orbit.points[15]) + + #async def _input_loop(self): + # fd = sys.stdin.fileno() + # old_settings = termios.tcgetattr(fd) + # try: + # tty.setcbreak(fd) + # while self._running: + # rlist, _, _ = select.select([sys.stdin], [], [], 0) + # if rlist: + # c = sys.stdin.read(1) + # if c == "q": # quit on q + # self._running = False + # else: + # await self._handle_key(c) + # await asyncio.sleep(0.001) + # finally: + # termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) async def _render_loop(self): import time @@ -93,22 +114,23 @@ class TerminalPlotter(pv.Plotter): np_image = numpy_support.vtk_to_numpy(vtk_array).reshape(height, width, 4) np_image = np_image[::-1] # Flip vertically np_image = np.ascontiguousarray(np_image) # Ensure memory layout is C-contiguous - print("Render & copy:", time.perf_counter() - start) + #print("Render & copy:", time.perf_counter() - start) - #await kitty.draw_to_terminal(np_image, width, height, compress=True) - await kitty.draw_to_terminal(np_image, width, height, use_shm=True) - print("Draw:", time.perf_counter() - start) - kitty.set_position(self.start_y, self.start_x) + #await self.terminal.draw_to_terminal(np_image, width, height, compress=True) + await self.terminal.draw(np_image, width, height, use_shm=True) + #print("Draw:", time.perf_counter() - start) + self.terminal.set_position(self.start_y, self.start_x) self.camera.Azimuth(1) async def run(self): await self.initialize() self.iren.initialize() - tasks = [ - asyncio.create_task(self._input_loop()), - asyncio.create_task(self._render_loop()), - ] - await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + #tasks = [ + # #asyncio.create_task(self._input_loop()), + # asyncio.create_task(self._render_loop()), + #] + #await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + await self._render_loop() self._running = False async def main():