kglobe/terminalplotter.py
2025-07-21 23:24:55 +02:00

124 lines
4.2 KiB
Python

import asyncio
import io
import math
import select
import sys
import termios
import tty
import pyvista as pv
import kitty
class TerminalPlotter(pv.Plotter):
def __init__(self, width, height, **kwargs):
super().__init__(off_screen=True, window_size=(width, height), **kwargs)
self.width = width
self.height = height
self._running = True
self._needs_render = True
self._render_complete = False
self._rendering = False
self._render_pending = True
self.set_background([0.0, 0.0, 0.0])
self.add_on_render_callback(self._on_render)
# if we are too close to the bottom of the terminal, create some space.
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()
self.needed_lines = math.ceil(self.height / h_pix)
# Add vertical space if needed
if self.rows - self.start_y < self.needed_lines:
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)
async def render_to_kitty(self):
self.render()
buf = io.BytesIO()
self.screenshot(buf, transparent_background=True, window_size=(self.width, self.height))
await kitty.draw_to_terminal(buf)
# print("y:", self.start_y, "rows:", self.rows, end="")
kitty.set_position(self.start_y, self.start_x)
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.01)
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)
self._needs_render = True
# await self.render_to_kitty()
def _on_render(self, *args):
"""Callback that runs after each render."""
self._rendering = False
self._render_complete = True
async def _display_loop(self):
"""Loop that captures the screen and updates the terminal after rendering."""
while self._running:
if self._render_complete:
self._render_complete = False
buf = io.BytesIO()
# Safe: Screenshot only after render is confirmed complete
self.screenshot(buf, transparent_background=True, window_size=(self.width, self.height))
await kitty.draw_to_terminal(buf)
kitty.set_position(self.start_y, self.start_x)
await asyncio.sleep(0.001)
async def _render_loop(self):
while self._running:
if self._needs_render and not self._rendering:
self._rendering = True
self._needs_render = False
self.update() # triggers a render and calls _on_render
await asyncio.sleep(0.001)
async def _display_in_terminal(self, buf):
await kitty.draw_to_terminal(buf)
kitty.set_position(self.start_y, self.start_x)
async def run(self):
await self.initialize()
self.iren.initialize()
self._needs_render = True
tasks = [
asyncio.create_task(self._input_loop()),
asyncio.create_task(self._render_loop()),
asyncio.create_task(self._display_loop()),
]
await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
self._running = False
#await self.render_to_kitty()
#await self._input_loop()