vimol 0.1.0__py3-none-any.whl

vimol/input.py

@@ -0,0 +1,252 @@
+"""Terminal input decoding: keys and SGR mouse events (cell *and* pixel).
+
+The decoder is stateful so it tolerates escape sequences that get split across
+reads. It is deliberately independent of the viewer: an embedding application
+can create an :class:`InputDecoder`, feed it whatever bytes it intercepts, and
+forward the resulting :class:`MouseEvent` / :class:`KeyEvent` objects to a
+:class:`vimol.widget.MoleculeWidget` — so you can capture the mouse inside
+your own UI region and drive the molecule with it.
+
+Mouse coordinates come back in *pixels* when the terminal supports SGR-Pixels
+(DECSET 1016, a Kitty extension also in Ghostty/WezTerm/foot) and in character
+cells otherwise; the ``pixel`` flag on each event says which.
+"""
+from __future__ import annotations
+
+import os
+import select
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+
+# --------------------------------------------------------------------------
+# Events
+# --------------------------------------------------------------------------
+@dataclass
+class MouseEvent:
+    action: str          # 'down' | 'up' | 'drag' | 'move' | 'scroll'
+    x: float             # pixels if pixel else 0-based cell column
+    y: float             # pixels if pixel else 0-based cell row
+    button: Optional[int] = None    # 0 left, 1 middle, 2 right (None for move/scroll)
+    scroll: Optional[str] = None     # 'up' | 'down' | 'left' | 'right'
+    shift: bool = False
+    alt: bool = False
+    ctrl: bool = False
+    pixel: bool = False   # True -> x/y are pixels, False -> x/y are cells
+
+
+@dataclass
+class KeyEvent:
+    key: str             # a single char, or a name: 'up' 'down' 'left' 'right'
+                         # 'escape' 'enter' 'tab' 'backspace'
+    def __str__(self) -> str:
+        return self.key
+
+
+Event = Union[MouseEvent, KeyEvent]
+
+_CSI_FINAL = {"A": "up", "B": "down", "C": "right", "D": "left",
+              "H": "home", "F": "end"}
+
+
+# --------------------------------------------------------------------------
+# Mouse enable / disable control sequences
+# --------------------------------------------------------------------------
+def enable_mouse(pixel: bool = True, hover: bool = False) -> bytes:
+    """Bytes to turn on SGR mouse reporting.
+
+    hover=True uses any-event tracking (1003) so motion is reported even with no
+    button held (needed for hover picking); otherwise button-event tracking
+    (1002) reports motion only while dragging. pixel=True additionally requests
+    SGR-Pixels (1016) for pixel-precise coordinates.
+    """
+    motion = "1003" if hover else "1002"
+    seq = f"\x1b[?1000;{motion};1006h"
+    if pixel:
+        seq += "\x1b[?1016h"
+    return seq.encode()
+
+
+def disable_mouse(pixel: bool = True) -> bytes:
+    seq = "\x1b[?1000;1002;1003;1006l"
+    if pixel:
+        seq += "\x1b[?1016l"
+    return seq.encode()
+
+
+def query_decset(mode: int, fd_in: int = 0, fd_out: int = 1, timeout: float = 0.2):
+    """Probe a DECSET private mode via DECRQM; return its reported value.
+
+    Returns the numeric value from the terminal's ``CSI ? mode ; value $y``
+    reply (1=set, 2=reset, 3=perm-set, 4=perm-reset, 0=not recognized) or None
+    if the terminal did not answer. Requires the tty to be in raw mode.
+    """
+    try:
+        os.write(fd_out, f"\x1b[?{mode}$p".encode())
+    except OSError:
+        return None
+    buf = b""
+    end_marker = b"$y"
+    import time as _time
+    # We cannot call time.time() indirectly? it's fine here (not a workflow script).
+    deadline = _time.monotonic() + timeout
+    while _time.monotonic() < deadline:
+        r, _, _ = select.select([fd_in], [], [], max(0.0, deadline - _time.monotonic()))
+        if not r:
+            break
+        try:
+            chunk = os.read(fd_in, 256)
+        except OSError:
+            break
+        if not chunk:
+            break
+        buf += chunk
+        if end_marker in buf:
+            break
+    # parse: ESC [ ? <mode> ; <value> $ y
+    marker = f"\x1b[?{mode};".encode()
+    i = buf.find(marker)
+    if i < 0:
+        return None
+    j = buf.find(end_marker, i)
+    if j < 0:
+        return None
+    try:
+        val = int(buf[i + len(marker):j])
+    except ValueError:
+        return None
+    return val
+
+
+def supports_pixel_mouse(fd_in: int = 0, fd_out: int = 1, timeout: float = 0.5) -> bool:
+    """True if the terminal reports SGR-Pixels (1016) as a recognized mode.
+
+    The default timeout is generous (0.5s) on purpose: the DECRQM round-trip
+    is one-time startup cost, but over SSH a tight timeout loses the race and
+    reports "no pixel mouse" for a terminal that in fact supports it. A wrong
+    answer here no longer breaks the mouse -- the viewer keeps the wire format
+    and the decoder consistent either way (see Viewer._setup_mouse) -- but a
+    right answer preserves pixel-precise dragging and picking over SSH.
+    """
+    val = query_decset(1016, fd_in, fd_out, timeout=timeout)
+    return val in (1, 2, 3, 4)
+
+
+# --------------------------------------------------------------------------
+# Decoder
+# --------------------------------------------------------------------------
+class InputDecoder:
+    """Turn raw terminal bytes into :class:`Event` objects.
+
+    Parameters
+    ----------
+    pixel:
+        Whether SGR mouse coordinates should be interpreted as pixels. Set this
+        to match whether you enabled mode 1016 (see :func:`supports_pixel_mouse`).
+    """
+
+    def __init__(self, pixel: bool = False):
+        self.pixel = pixel
+        self._buf = bytearray()
+
+    def feed(self, data: bytes) -> List[Event]:
+        self._buf += data
+        events: List[Event] = []
+        while self._buf:
+            ev, consumed = self._parse_one()
+            if consumed == 0:
+                break  # incomplete sequence, wait for more bytes
+            if ev is not None:
+                events.append(ev)
+            del self._buf[:consumed]
+        return events
+
+    def flush(self) -> List[Event]:
+        """Emit any pending unterminated sequence (e.g. a lone ESC -> escape)."""
+        if len(self._buf) == 1 and self._buf[0] == 0x1B:
+            self._buf.clear()
+            return [KeyEvent("escape")]
+        return []
+
+    # -- internals --------------------------------------------------------
+    def _parse_one(self):
+        b = self._buf
+        c = b[0]
+        if c != 0x1B:
+            return self._parse_key_byte(c), 1
+        # escape sequence
+        if len(b) < 2:
+            return None, 0  # wait; flush() handles a lingering lone ESC
+        if b[1] == ord("["):
+            return self._parse_csi()
+        if b[1] == ord("O"):  # SS3 (application arrows)
+            if len(b) < 3:
+                return None, 0
+            name = _CSI_FINAL.get(chr(b[2]))
+            return (KeyEvent(name) if name else None), 3
+        # Alt+key
+        return KeyEvent(chr(b[1])), 2
+
+    def _parse_key_byte(self, c: int):
+        if c == 0x0D:
+            return KeyEvent("enter")
+        if c == 0x09:
+            return KeyEvent("tab")
+        if c in (0x7F, 0x08):
+            return KeyEvent("backspace")
+        return KeyEvent(chr(c))
+
+    def _parse_csi(self):
+        b = self._buf
+        n = len(b)
+        # SGR mouse: ESC [ < params (M|m)
+        if n >= 3 and b[2] == ord("<"):
+            k = 3
+            while k < n and b[k] not in (ord("M"), ord("m")):
+                k += 1
+            if k >= n:
+                return None, 0  # incomplete
+            body = bytes(b[3:k]).decode("ascii", "ignore")
+            final = chr(b[k])
+            return self._decode_sgr_mouse(body, final), k + 1
+        # generic CSI: read until final byte in @..~
+        k = 2
+        while k < n and not (0x40 <= b[k] <= 0x7E):
+            k += 1
+        if k >= n:
+            return None, 0
+        final = chr(b[k])
+        name = _CSI_FINAL.get(final)
+        return (KeyEvent(name) if name else None), k + 1
+
+    def _decode_sgr_mouse(self, body: str, final: str) -> Optional[MouseEvent]:
+        try:
+            bs, xs, ys = body.split(";")
+            code = int(bs); x = int(xs); y = int(ys)
+        except ValueError:
+            return None
+        shift = bool(code & 4)
+        alt = bool(code & 8)
+        ctrl = bool(code & 16)
+        wheel = bool(code & 64)
+        motion = bool(code & 32)
+        low = code & 3
+        # coordinates: SGR is 1-based; make cells 0-based. Pixels stay as-is.
+        if self.pixel:
+            px, py = float(x), float(y)
+        else:
+            px, py = float(x - 1), float(y - 1)
+
+        if wheel:
+            direction = {0: "up", 1: "down", 2: "left", 3: "right"}.get(low, "up")
+            return MouseEvent("scroll", px, py, scroll=direction,
+                              shift=shift, alt=alt, ctrl=ctrl, pixel=self.pixel)
+        if motion:
+            if low == 3:
+                return MouseEvent("move", px, py, button=None,
+                                  shift=shift, alt=alt, ctrl=ctrl, pixel=self.pixel)
+            return MouseEvent("drag", px, py, button=low,
+                              shift=shift, alt=alt, ctrl=ctrl, pixel=self.pixel)
+        action = "down" if final == "M" else "up"
+        return MouseEvent(action, px, py, button=low,
+                          shift=shift, alt=alt, ctrl=ctrl, pixel=self.pixel)

vimol/kitty.py

@@ -0,0 +1,297 @@
+"""Kitty terminal graphics protocol encoder + terminal geometry helpers.
+
+Reference: https://sw.kovidgoyal.net/kitty/graphics-protocol/
+
+We transmit raw RGB/RGBA pixels (optionally zlib-compressed) in <=4096-byte
+base64 chunks. The public surface is small on purpose so the module can be
+lifted into other terminal apps unchanged.
+"""
+from __future__ import annotations
+
+import base64
+import os
+import sys
+import zlib
+from typing import Optional, Tuple
+
+import numpy as np
+
+_ESC = b"\x1b"
+_GRAPHICS_START = b"\x1b_G"
+_GRAPHICS_END = b"\x1b\\"
+_CHUNK = 4096
+
+
+# --------------------------------------------------------------------------
+# Terminal geometry
+# --------------------------------------------------------------------------
+def terminal_size_px(fd: int = 1) -> Tuple[int, int, int, int]:
+    """Return (cols, rows, width_px, height_px) for the terminal on *fd*.
+
+    Falls back to environment / sane defaults when the ioctl is unavailable
+    (e.g. output is not a tty). width/height px may be 0 if the terminal does
+    not report pixel dimensions.
+    """
+    cols, rows, xpx, ypx = 80, 24, 0, 0
+    try:
+        import fcntl
+        import struct
+        import termios
+
+        buf = fcntl.ioctl(fd, termios.TIOCGWINSZ, b"\x00" * 8)
+        rows, cols, xpx, ypx = struct.unpack("HHHH", buf)
+    except Exception:
+        cols = int(os.environ.get("COLUMNS", cols))
+        rows = int(os.environ.get("LINES", rows))
+    return cols, rows, xpx, ypx
+
+
+def cell_size_px(fd: int = 1) -> Tuple[float, float]:
+    """(cell_width_px, cell_height_px). Defaults to 9x18 if unknown.
+
+    Derived by dividing the window's reported pixel extent by its cell
+    count. That extent can include the terminal's window padding, which
+    this then smears across every row/column as a fractional per-cell
+    error -- invisible to the continuous, radius-based atom picker but
+    enough to make a rigid per-cell grid (the periodic-table picker) land
+    on the wrong cell, worse the further down you go as the error adds up.
+    Rounding removes the smaller part of that noise; :func:`query_cell_size_px`
+    removes it entirely by asking the terminal for the exact cell size, and
+    the viewer prefers that when the terminal answers.
+    """
+    cols, rows, xpx, ypx = terminal_size_px(fd)
+    cw = round(xpx / cols) if xpx and cols else 9.0
+    ch = round(ypx / rows) if ypx and rows else 18.0
+    return float(cw), float(ch)
+
+
+def query_cell_size_px(fd_in: int = 0, fd_out: int = 1, timeout: float = 0.2):
+    """Ask the terminal for its exact cell size in pixels via ``CSI 16 t``.
+
+    Returns ``(cell_width_px, cell_height_px)`` from the terminal's
+    ``CSI 6 ; height ; width t`` reply, or ``None`` if it doesn't answer.
+    Unlike :func:`cell_size_px` this is the terminal's own authoritative
+    cell metric, with no window-padding contamination -- so cell/pixel
+    hit-testing lines up exactly with where glyphs are actually drawn.
+    Requires the tty to be in raw mode (mirrors input.query_decset).
+    """
+    import select
+    import time as _time
+
+    try:
+        os.write(fd_out, b"\x1b[16t")
+    except OSError:
+        return None
+    buf = b""
+    deadline = _time.monotonic() + timeout
+    while _time.monotonic() < deadline:
+        r, _, _ = select.select([fd_in], [], [], max(0.0, deadline - _time.monotonic()))
+        if not r:
+            break
+        try:
+            chunk = os.read(fd_in, 64)
+        except OSError:
+            break
+        if not chunk:
+            break
+        buf += chunk
+        if b"t" in buf and b"\x1b[6;" in buf:
+            break
+    # parse: ESC [ 6 ; <height> ; <width> t
+    marker = b"\x1b[6;"
+    i = buf.find(marker)
+    if i < 0:
+        return None
+    j = buf.find(b"t", i)
+    if j < 0:
+        return None
+    try:
+        h_str, w_str = buf[i + len(marker):j].split(b";")
+        ch = float(int(h_str))
+        cw = float(int(w_str))
+    except (ValueError, IndexError):
+        return None
+    if cw <= 0 or ch <= 0:
+        return None
+    return cw, ch
+
+
+def unique_id_base(stride: int = 4) -> int:
+    """A per-process base for Kitty graphics ``i=``/``p=`` ids.
+
+    Image ids live in a namespace global to the whole kitty *process* (shared
+    across all its panes/tabs), not per-window — see the graphics protocol
+    spec's note that "IDs are in a global namespace [so] there can easily be
+    collisions." Hardcoding small ids like 1/2 means two independent
+    vimol instances sharing a kitty process (e.g. two panes) can delete or
+    overwrite each other's frames. Deriving the base from the pid keeps
+    concurrent instances apart; the stride ensures that even OS-assigned
+    sequential pids (common right after spawning several processes) don't
+    produce adjacent, overlapping id ranges.
+    """
+    pid = os.getpid()
+    return ((pid * stride) % 0x7FFFFFFF) or 1
+
+
+def supports_kitty() -> bool:
+    """Best-effort detection of Kitty graphics support via environment."""
+    if os.environ.get("VIMOL_FORCE_KITTY"):
+        return True
+    if os.environ.get("KITTY_WINDOW_ID"):
+        return True
+    term = os.environ.get("TERM", "")
+    if "kitty" in term or "ghostty" in term:
+        return True
+    prog = os.environ.get("TERM_PROGRAM", "").lower()
+    if prog in ("ghostty", "wezterm"):
+        return True
+    if os.environ.get("WEZTERM_PANE"):
+        return True
+    return False
+
+
+# --------------------------------------------------------------------------
+# Image encoding
+# --------------------------------------------------------------------------
+def _controls(d: dict) -> bytes:
+    return ",".join(f"{k}={v}" for k, v in d.items()).encode("ascii")
+
+
+def encode_image(
+    pixels: np.ndarray,
+    *,
+    image_id: int = 1,
+    placement_id: Optional[int] = None,
+    cols: Optional[int] = None,
+    rows: Optional[int] = None,
+    move_cursor: bool = False,
+    compress: bool = True,
+    compress_level: int = 6,
+    z_index: int = 0,
+    quiet: int = 2,
+) -> bytes:
+    """Encode an (H, W, 3|4) uint8 array as Kitty graphics-protocol bytes.
+
+    cols/rows scale the image into that many terminal cells (defaults to the
+    image's native pixel size). When *move_cursor* is False (C=1) the cursor is
+    left in place, which is what you want when compositing a UI around the
+    image. placement_id defaults to image_id (placements are scoped to their
+    image, so this is always collision-safe).
+
+    *compress_level* is the zlib level (0-9). The compress step costs about as
+    much per frame as the whole raycast, so interactive callers pass level 1
+    (roughly half the CPU of the default 6 for ~25% more bytes -- a fine
+    trade for a local terminal); a resting still can afford the default for a
+    smaller payload.
+    """
+    if placement_id is None:
+        placement_id = image_id
+    arr = np.ascontiguousarray(pixels, dtype=np.uint8)
+    h, w = arr.shape[0], arr.shape[1]
+    fmt = 32 if arr.ndim == 3 and arr.shape[2] == 4 else 24
+    raw = arr.tobytes()
+    if compress:
+        raw = zlib.compress(raw, compress_level)
+    payload = base64.standard_b64encode(raw)
+
+    ctrl = {
+        "a": "T",          # transmit and display
+        "f": fmt,          # 24=RGB, 32=RGBA
+        "s": w,
+        "v": h,
+        "i": image_id,
+        "p": placement_id,
+        "q": quiet,        # 0=verbose, 1=no ok, 2=no ok/err
+    }
+    if compress:
+        ctrl["o"] = "z"
+    if cols:
+        ctrl["c"] = int(cols)
+    if rows:
+        ctrl["r"] = int(rows)
+    if not move_cursor:
+        ctrl["C"] = 1
+    if z_index:
+        ctrl["z"] = int(z_index)
+
+    out = bytearray()
+    # chunked transfer: first chunk carries controls, subsequent carry only m=
+    if len(payload) <= _CHUNK:
+        ctrl["m"] = 0
+        out += _GRAPHICS_START + _controls(ctrl) + b";" + payload + _GRAPHICS_END
+        return bytes(out)
+
+    first = True
+    view = memoryview(payload)
+    n = len(payload)
+    pos = 0
+    while pos < n:
+        chunk = view[pos:pos + _CHUNK]
+        pos += _CHUNK
+        last = pos >= n
+        if first:
+            c = dict(ctrl)
+            c["m"] = 0 if last else 1
+            out += _GRAPHICS_START + _controls(c) + b";" + bytes(chunk) + _GRAPHICS_END
+            first = False
+        else:
+            m = 0 if last else 1
+            out += _GRAPHICS_START + f"m={m}".encode() + b";" + bytes(chunk) + _GRAPHICS_END
+    return bytes(out)
+
+
+def delete_image(image_id: int = 1) -> bytes:
+    """Bytes to delete a transmitted image (and its placements) by id."""
+    return _GRAPHICS_START + _controls({"a": "d", "d": "I", "i": image_id, "q": 2}) + b";" + _GRAPHICS_END
+
+
+def clear_all_images() -> bytes:
+    return _GRAPHICS_START + b"a=d,d=A,q=2;" + _GRAPHICS_END
+
+
+def set_pointer_shape(shape: str) -> bytes:
+    """OSC 22 bytes to PUSH the OS mouse-pointer icon to a CSS-cursor *shape*.
+
+    Kitty-family terminals honor this while focused; others ignore an OSC they
+    don't recognize, so it's safe to write unconditionally (like the mouse-mode
+    sequences). Pushing (rather than a plain set) pairs with
+    :func:`reset_pointer_shape`'s pop, so disarming restores whatever shape
+    was actually active before -- not just the terminal's generic default,
+    which a plain set/reset pair would revert to instead. Returns raw bytes
+    for the caller to write.
+    """
+    return f"\x1b]22;>{shape}\x1b\\".encode()
+
+
+def reset_pointer_shape() -> bytes:
+    """OSC 22 bytes to POP the pointer shape pushed by :func:`set_pointer_shape`."""
+    return b"\x1b]22;<\x1b\\"
+
+
+def write_bytes(data: bytes, fd: int = 1) -> None:
+    os.write(fd, data)
+
+
+def png_bytes(pixels: np.ndarray) -> bytes:
+    """Encode an (H, W, 3|4) uint8 array as a PNG (stdlib zlib, no Pillow)."""
+    import struct
+
+    arr = np.ascontiguousarray(pixels, np.uint8)
+    h, w = arr.shape[0], arr.shape[1]
+    channels = arr.shape[2] if arr.ndim == 3 else 1
+    color_type = {1: 0, 3: 2, 4: 6}[channels]
+
+    def chunk(tag: bytes, data: bytes) -> bytes:
+        return struct.pack(">I", len(data)) + tag + data + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)
+
+    sig = b"\x89PNG\r\n\x1a\n"
+    ihdr = struct.pack(">IIBBBBB", w, h, 8, color_type, 0, 0, 0)
+    # add filter byte 0 per scanline
+    stride = w * channels
+    raw = bytearray()
+    flat = arr.reshape(h, stride)
+    for y in range(h):
+        raw.append(0)
+        raw += flat[y].tobytes()
+    idat = zlib.compress(bytes(raw), 9)
+    return sig + chunk(b"IHDR", ihdr) + chunk(b"IDAT", idat) + chunk(b"IEND", b"")