Add a screenshot remote control command

docs/changelog.rst

@@ -179,6 +179,12 @@ consumption to do the same tasks.
 Detailed list of changes
 -------------------------------------
 
+0.49.0 [future]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Add a new :code:`kitten @ screenshot` remote control command to take a pixel perfect PNG screenshot of an OS Window, tab or window
+
+
 0.48.0 [future]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

kitty/boss.py

@@ -9,8 +9,10 @@ import re
 import socket
 import subprocess
 import sys
+from collections import deque
 from collections.abc import Callable, Container, Generator, Iterable, Iterator, Sequence
 from contextlib import contextmanager, suppress
+from dataclasses import dataclass
 from functools import partial
 from gettext import gettext as _
 from gettext import ngettext
@@ -110,6 +112,7 @@ from .fast_data_types import (
     os_window_focus_counters,
     os_window_font_size,
     redirect_mouse_handling,
+    request_callback_with_thumbnail,
     ring_bell,
     run_with_activation_token,
     safe_pipe,
@@ -186,6 +189,19 @@ if TYPE_CHECKING:
 
 RCResponse = Union[dict[str, Any], None, AsyncResponse]
 
+ThumbnailCallback = Callable[[int, int, bytes, int, int], None]
+
+
+@dataclass
+class PendingThumbnailRequest:
+    os_window_id: int
+    callback: ThumbnailCallback
+    window_id: int = 0
+    include_tab_bar: bool = False
+    scale: float = 0.25
+    max_width: int = 480
+    no_scaling: bool = False
+
 
 class OSWindowDict(TypedDict):
     id: int
@@ -413,6 +429,7 @@ class Boss:
         self.cached_values = cached_values
         self.os_window_map: dict[int, TabManager] = {}
         self.os_window_death_actions: dict[int, Callable[[], None]] = {}
+        self.thumbnail_request_queue: deque[PendingThumbnailRequest] = deque()
         self.cursor_blinking = True
         self.shutting_down = False
         self.misc_config_errors: list[str] = []
@@ -1436,6 +1453,43 @@ class Boss:
         if tm := self.os_window_map.get(os_window_id):
             tm.start_window_drag(pixels, width, height)
 
+    def request_thumbnail(
+        self, os_window_id: int, callback: ThumbnailCallback, window_id: int = 0, include_tab_bar: bool = False,
+        scale: float = 0.25, max_width: int = 480, no_scaling: bool = False,
+    ) -> None:
+        # request_callback_with_thumbnail() only supports a single request being in flight at
+        # a time (it uses a single global slot in the C code), so serialize all
+        # requests (window/tab drag thumbnails as well as screenshot RC command
+        # requests) through this queue.
+        req = PendingThumbnailRequest(os_window_id, callback, window_id, include_tab_bar, scale, max_width, no_scaling)
+        was_idle = not self.thumbnail_request_queue
+        self.thumbnail_request_queue.append(req)
+        if was_idle:
+            self._issue_next_thumbnail_request()
+
+    def _issue_next_thumbnail_request(self) -> None:
+        while self.thumbnail_request_queue:
+            req = self.thumbnail_request_queue[0]
+            if req.os_window_id not in self.os_window_map:
+                # the OS window went away while this request was queued, fail it and move on
+                # rather than leaving every request behind it stuck forever
+                self.thumbnail_request_queue.popleft()
+                req.callback(req.os_window_id, req.window_id, b'', 0, 0)
+                continue
+            request_callback_with_thumbnail(
+                'thumbnail_ready', req.os_window_id, req.window_id, req.include_tab_bar,
+                req.scale, req.max_width, req.no_scaling)
+            break
+
+    def thumbnail_ready(self, os_window_id: int, window_id: int, pixels: bytes, width: int, height: int) -> None:
+        if not self.thumbnail_request_queue:
+            return
+        req = self.thumbnail_request_queue.popleft()
+        try:
+            req.callback(os_window_id, window_id, pixels, width, height)
+        finally:
+            self._issue_next_thumbnail_request()
+
     def on_window_resize(self, os_window_id: int, w: int, h: int, dpi_changed: bool) -> None:
         if dpi_changed:
             self.on_dpi_change(os_window_id)

kitty/child-monitor.c

@@ -878,16 +878,20 @@ thumbnail_callback(OSWindow *os_window) {
         }
     }
     unsigned vw = region.right - region.left, vh = region.bottom - region.top;
-    unsigned thumb_w = (unsigned)(vw * tc.scale), thumb_h = (unsigned)(vh * tc.scale);
-    if (thumb_w > tc.max_width) {
-        thumb_w = tc.max_width;
-        double scale = 300. / vw;
-        thumb_h = (unsigned)(vh * scale + 0.5f);
+    unsigned thumb_w, thumb_h;
+    if (tc.no_scaling) { thumb_w = vw; thumb_h = vh; }
+    else {
+        thumb_w = (unsigned)(vw * tc.scale); thumb_h = (unsigned)(vh * tc.scale);
+        if (thumb_w > tc.max_width) {
+            thumb_w = tc.max_width;
+            double scale = 300. / vw;
+            thumb_h = (unsigned)(vh * scale + 0.5f);
+        }
     }
     RAII_PyObject(pixels, PyBytes_FromStringAndSize(NULL, (Py_ssize_t)4 * thumb_w * thumb_h));
     if (pixels && global_state.boss) {
         take_screenshot_of_rectangular_region(
-            os_window, region, (unsigned char*)PyBytes_AS_STRING(pixels), &thumb_w, &thumb_h);
+            os_window, region, (unsigned char*)PyBytes_AS_STRING(pixels), &thumb_w, &thumb_h, tc.no_scaling);
         _PyBytes_Resize(&pixels, (Py_ssize_t)4 * thumb_w * thumb_h);
         PyObject *r = PyObject_CallMethod(
             global_state.boss, tc.callback, "KKOII", os_window->id, tc.window, pixels, thumb_w, thumb_h);

kitty/fast_data_types.pyi

@@ -1854,7 +1854,7 @@ def set_window_being_dragged(window_id: int = 0, drag_started: bool = False, x:
 def get_window_being_dragged() -> tuple[int, bool, float, float]: ...
 def request_callback_with_thumbnail(
         callback: str, os_window_id: int, window_id: int = 0, include_tab_bar: bool = False,
-        scale: float = 0.25, max_width: int = 480
+        scale: float = 0.25, max_width: int = 480, no_scaling: bool = False
 ) -> None: ...
 def png_from_32bit_rgba_data(data: bytes, width: int, height: int, flip_vertically: bool = False) -> bytes: ...
 def set_uint_at_address(address: int, value: int) -> None: ...

kitty/rc/screenshot.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python
+# License: GPLv3 Copyright: 2024, Kovid Goyal <kovid at kovidgoyal.net>
+
+import os
+from base64 import standard_b64encode
+from typing import TYPE_CHECKING
+
+from kitty.fast_data_types import png_from_32bit_rgba_data
+from kitty.types import AsyncResponse
+
+from .base import (
+    MATCH_TAB_OPTION,
+    MATCH_WINDOW_OPTION,
+    ArgsType,
+    Boss,
+    MatchError,
+    PayloadGetType,
+    PayloadType,
+    RCOptions,
+    RemoteCommand,
+    RemoteControlErrorWithoutTraceback,
+    ResponseType,
+    Window,
+)
+
+if TYPE_CHECKING:
+    from kitty.cli_stub import ScreenshotRCOptions as CLIOptions
+
+
+class Screenshot(RemoteCommand):
+
+    protocol_spec = __doc__ = '''
+    match/str: The window to screenshot
+    match_tab/str: The tab to screenshot
+    output_path/str: Path to save the PNG image to, on the computer kitty is running on. Empty to return the image data instead.
+    '''
+
+    short_desc = 'Take a screenshot of a kitty OS window, tab or window'
+    desc = (
+        'Take a screenshot, as a PNG image, of an entire kitty OS window. Restrict the screenshot to a single'
+        ' kitty window or tab with :option:`--match`/:option:`--match-tab`. The specified window/tab must be'
+        ' currently visible, i.e. it must be in the active tab of its OS window and, in the case of a window,'
+        ' not hidden behind another window by the layout, otherwise this command will fail.\n\n'
+
+        'By default, the PNG image data is written to STDOUT. If instead a file path is specified, kitty itself'
+        ' (rather than this :program:`kitten` process) saves the screenshot to that path. Since kitty and the'
+        ' kitten are typically running on the same computer, this avoids copying the (potentially large)'
+        ' image data over the possibly slow remote control transport.'
+    )
+    options_spec = MATCH_WINDOW_OPTION + '\n\n' + MATCH_TAB_OPTION.replace('--match -m', '--match-tab -t')
+    args = RemoteCommand.Args(
+        spec='[OUTPUT_FILE]', json_field='output_path',
+        special_parse='!read_screenshot_args(io_data, args)',
+        completion=RemoteCommand.CompletionSpec.from_string('type:file ext:png'),
+    )
+    is_asynchronous = True
+
+    def message_to_kitty(self, global_opts: RCOptions, opts: 'CLIOptions', args: ArgsType) -> PayloadType:
+        if len(args) > 1:
+            self.fatal('Must specify at most one output file')
+        return {
+            'match': opts.match, 'match_tab': opts.match_tab,
+            'output_path': args[0] if args else '',
+        }
+
+    def response_from_kitty(self, boss: Boss, window: Window | None, payload_get: PayloadGetType) -> ResponseType:
+        match = payload_get('match')
+        match_tab = payload_get('match_tab')
+        target_window_id = 0
+        include_tab_bar = False
+
+        if match:
+            windows = list(boss.match_windows(match, window))
+            if not windows:
+                raise MatchError(match)
+            w = windows[0]
+            tab = w.tabref()
+            tm = tab.tab_manager_ref() if tab is not None else None
+            if tab is None or tm is None or tm.active_tab is not tab or not w.is_visible_in_layout:
+                raise RemoteControlErrorWithoutTraceback(
+                    'The matched window is not currently visible, screenshots can only be taken of visible windows')
+            os_window_id = w.os_window_id
+            target_window_id = w.id
+        elif match_tab:
+            tabs = list(boss.match_tabs(match_tab))
+            if not tabs:
+                raise MatchError(match_tab, 'tabs')
+            tab = tabs[0]
+            tm = tab.tab_manager_ref()
+            if tm is None or tm.active_tab is not tab:
+                raise RemoteControlErrorWithoutTraceback(
+                    'The matched tab is not currently visible, screenshots can only be taken of visible tabs')
+            os_window_id = tab.os_window_id
+        else:
+            atm = boss.active_tab_manager
+            if atm is None:
+                raise RemoteControlErrorWithoutTraceback('There is no active OS window to screenshot')
+            os_window_id = atm.os_window_id
+            include_tab_bar = True
+
+        output_path = payload_get('output_path') or ''
+        responder = self.create_async_responder(payload_get, window)
+
+        def callback(cb_os_window_id: int, cb_window_id: int, pixels: bytes, width: int, height: int) -> None:
+            if not pixels:
+                responder.send_error('Failed to take screenshot, the OS window may have been closed')
+                return
+            try:
+                png_data = png_from_32bit_rgba_data(pixels, width, height, True)
+                if output_path:
+                    with open(os.path.expanduser(output_path), 'wb') as f:
+                        f.write(png_data)
+                    responder.send_data(True)
+                else:
+                    responder.send_data(standard_b64encode(png_data).decode('ascii'))
+            except Exception as e:
+                responder.send_error(f'Failed to save screenshot: {e}')
+
+        boss.request_thumbnail(
+            os_window_id, callback, window_id=target_window_id, include_tab_bar=include_tab_bar, no_scaling=True)
+        return AsyncResponse()
+
+
+screenshot = Screenshot()

kitty/shaders.c

@@ -1682,15 +1682,27 @@ setup_os_window_for_rendering(OSWindow *os_window, Tab *tab, Window *active_wind
 // Scaling is performed on the GPU using the SCREENSHOT_PROGRAM shader for better performance.
 // The shader properly handles sRGB color space conversion and downscaling.
 // Setting the thumbnail dimensions to zero disables scaling.
+// If no_scaling is true, thumb_w/thumb_h are ignored and the region is read
+// back directly from the framebuffer with no GPU resampling, for a pixel
+// perfect capture (the output is in GL's native bottom-up row order, unlike
+// the scaled path below which flips to top-down).
 void
-take_screenshot_of_rectangular_region(OSWindow *os_window, Region region, unsigned char *dst_buf, unsigned *thumb_w, unsigned *thumb_h) {
-    unsigned vw = os_window->viewport_width;
+take_screenshot_of_rectangular_region(OSWindow *os_window, Region region, unsigned char *dst_buf, unsigned *thumb_w, unsigned *thumb_h, bool no_scaling) {
     unsigned vh = os_window->viewport_height;
 
     // Calculate the source region dimensions
     unsigned src_height = region.bottom - region.top;
     unsigned src_width = region.right - region.left;
 
+    if (no_scaling) {
+        *thumb_w = src_width; *thumb_h = src_height;
+        glPixelStorei(GL_PACK_ALIGNMENT, 1);
+        glReadPixels((GLint)region.left, (GLint)(vh - region.bottom), (GLsizei)src_width, (GLsizei)src_height, GL_RGBA, GL_UNSIGNED_BYTE, dst_buf);
+        glPixelStorei(GL_PACK_ALIGNMENT, 4);
+        return;
+    }
+
+    unsigned vw = os_window->viewport_width;
     if (!*thumb_w) *thumb_w = src_width;
     if (!*thumb_h) *thumb_h = src_height;
     *thumb_w = MIN(src_width, *thumb_w);

kitty/state.c

@@ -1751,9 +1751,9 @@ get_window_being_dragged(PyObject *self UNUSED, PyObject *args UNUSED) {
 static PyObject*
 request_callback_with_thumbnail(PyObject *self UNUSED, PyObject *args) {
     unsigned long long os_window_id, window_id = 0;
-    const char *callback; int include_tab_bar = 0;
+    const char *callback; int include_tab_bar = 0, no_scaling = 0;
     double scale = 0.25; unsigned max_width = 480;
-    if (!PyArg_ParseTuple(args, "sK|KpdI", &callback, &os_window_id, &window_id, &include_tab_bar, &scale, &max_width)) return NULL;
+    if (!PyArg_ParseTuple(args, "sK|KpdIp", &callback, &os_window_id, &window_id, &include_tab_bar, &scale, &max_width, &no_scaling)) return NULL;
     WITH_OS_WINDOW(os_window_id)
         global_state.thumbnail_callback.os_window = os_window->id;
         global_state.thumbnail_callback.window = window_id;
@@ -1761,6 +1761,7 @@ request_callback_with_thumbnail(PyObject *self UNUSED, PyObject *args) {
         snprintf(global_state.thumbnail_callback.callback, arraysz(global_state.thumbnail_callback.callback), "%s", callback);
         global_state.thumbnail_callback.max_width = max_width;
         global_state.thumbnail_callback.scale = scale;
+        global_state.thumbnail_callback.no_scaling = no_scaling;
         mark_os_window_dirty(os_window_id);
     END_WITH_OS_WINDOW
     Py_RETURN_NONE;

kitty/state.h

@@ -523,6 +523,7 @@ typedef struct GlobalState {
         char callback[32];
         bool include_tab_bar;
         double scale; unsigned max_width;
+        bool no_scaling;
     } thumbnail_callback;
     struct {
         id_type id; bool drag_started;
@@ -640,7 +641,7 @@ void dispatch_buffered_keys(Window *w);
 bool screen_needs_rendering_in_layers(OSWindow *os_window, Window *w, Screen *screen);
 void setup_os_window_for_rendering(OSWindow*, Tab*, Window*, bool);
 void swap_window_buffers(OSWindow *w);
-void take_screenshot_of_rectangular_region(OSWindow *os_window, Region region, unsigned char *dst_buf, unsigned *thumb_w, unsigned *thumb_h);
+void take_screenshot_of_rectangular_region(OSWindow *os_window, Region region, unsigned char *dst_buf, unsigned *thumb_w, unsigned *thumb_h, bool no_scaling);
 bool current_framebuffer_is_ok(void);
 void request_drop_status_update(OSWindow *osw);
 void register_mimes_for_drop(OSWindow *w, const char **mimes, size_t sz);

kitty/tabs.py

@@ -43,7 +43,6 @@ from .fast_data_types import (
     remove_window,
     reorder_tabs,
     replace_c0_codes_except_nl_space_tab,
-    request_callback_with_thumbnail,
     ring_bell,
     set_active_tab,
     set_active_window,
@@ -1803,7 +1802,7 @@ class TabManager:  # {{{
                 threshold = get_options().drag_threshold
                 if threshold and math.sqrt((x-start_x)**2 + (y-start_y)**2) > threshold:
                     set_tab_being_dragged(dragged_tab_id, True, start_x, start_y)
-                    request_callback_with_thumbnail("start_tab_drag", self.os_window_id)
+                    get_boss().request_thumbnail(self.os_window_id, get_boss().start_tab_drag)
                     self.recent_tab_bar_mouse_events.clear()
             return
 
@@ -1859,7 +1858,7 @@ class TabManager:  # {{{
                 dist_sq = (x - start_x)**2 + (y - start_y)**2
                 if threshold and dist_sq > threshold * threshold:
                     set_window_being_dragged(dragged_window_id, True, start_x, start_y)
-                    request_callback_with_thumbnail("start_window_drag", self.os_window_id, dragged_window_id)
+                    boss.request_thumbnail(self.os_window_id, boss.start_window_drag, window_id=dragged_window_id)
                     self.recent_title_bar_mouse_events.clear()
             return
         self.recent_title_bar_mouse_events.add(button, modifiers, action, x, y, window_id)

local-agent.md

@@ -45,11 +45,31 @@ To run a Go test named TestMyFunction, use:
 ./test.py MyFunction
 ```
 
+## Remote control API for verification
+
+kitty has a comprehensive remote control API you can use for manual verification of
+your changes. Run kitty as:
+
+    kitty -o allow_remote_control=y --listen-on=@test-kitty-xxx
+
+Then, you can take a screenshot of kitty and save it to test.png with:
+
+    kitten @ --to=@test-kitty-xxx screenshot test.png
+
+You can create window and tabs, send key events to kitty, query kitty
+state, etc using the various remote control sub-commands, which you can query
+using:
+
+    kitten @ --help
+
+
 ## Verification Pipeline
 
 Before declaring a task complete, you must follow this exact verification lifecycle:
 1. Run the local **Build Command** to guarantee zero compilation or compilation-stage type errors.
 2. Run the local **Test Command** to run the full test suite
 3. If errors occur, analyze the stdout logs completely before writing a fix. Do not guess.
+4. If your changes involve rendering changes to kitty manually verify
+   them by running kitty and using the remote control API as described above.
 4. If the change you have made is user facing, update the docs/changelog.rst
    file with a brief description of your changes

tools/cmd/at/screenshot.go

@@ -0,0 +1,39 @@
+// License: GPLv3 Copyright: 2024, Kovid Goyal, <kovid at kovidgoyal.net>
+
+package at
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/emmansun/base64"
+)
+
+func screenshot_handle_response(data []byte) error {
+	png_data, err := base64.StdEncoding.DecodeString(string(data))
+	if err != nil {
+		return err
+	}
+	_, err = os.Stdout.Write(png_data)
+	return err
+}
+
+func read_screenshot_args(io_data *rc_io_data, args []string) (func(io_data *rc_io_data) (bool, error), error) {
+	if len(args) > 1 {
+		return nil, fmt.Errorf("%s", "Must specify at most one output file")
+	}
+	if len(args) == 0 {
+		// kitty writes the screenshot directly to the output file when one is
+		// given (it runs on the same computer as kitty), so a custom response
+		// handler is only needed to write the PNG data to STDOUT.
+		io_data.handle_response = screenshot_handle_response
+	}
+	return func(io_data *rc_io_data) (bool, error) {
+		// io_data.rc.Payload is only populated after this generator is created,
+		// so the payload field must be set here rather than above.
+		if len(args) == 1 {
+			set_payload_string_field(io_data, "Output_path", args[0])
+		}
+		return true, nil
+	}, nil
+}