From 22f6728fed3b7e6f09525915613a171ddf280648 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Wed, 1 Mar 2023 17:34:38 +0530 Subject: [PATCH] Do not buffer PNG data to disk when setting window background or logo images --- kitty/boss.py | 4 ++-- kitty/fast_data_types.pyi | 2 +- kitty/rc/base.py | 33 ++++++++++++++------------------ kitty/rc/set_background_image.py | 9 ++++----- kitty/rc/set_window_logo.py | 14 +++++++------- kitty/state.c | 16 +++++++++------- kitty/window.py | 4 ++-- kitty/window_logo.c | 14 ++++++++++++-- kitty/window_logo.h | 2 +- 9 files changed, 52 insertions(+), 46 deletions(-) diff --git a/kitty/boss.py b/kitty/boss.py index 8ef34700e..cb1a4e40e 100644 --- a/kitty/boss.py +++ b/kitty/boss.py @@ -2639,8 +2639,8 @@ class Boss: self.choose_entry('Choose an OS window to move the tab to', items, chosen) - def set_background_image(self, path: Optional[str], os_windows: Tuple[int, ...], configured: bool, layout: Optional[str]) -> None: - set_background_image(path, os_windows, configured, layout) + def set_background_image(self, path: Optional[str], os_windows: Tuple[int, ...], configured: bool, layout: Optional[str], png_data: bytes = b'') -> None: + set_background_image(path, os_windows, configured, layout, png_data) for os_window_id in os_windows: self.default_bg_changed_for(os_window_id) diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi index fddc8c18a..62a2e8be9 100644 --- a/kitty/fast_data_types.pyi +++ b/kitty/fast_data_types.pyi @@ -1351,7 +1351,7 @@ def send_mouse_event(screen: Screen, x: int, y: int, button: int, action: int, m pass -def set_window_logo(os_window_id: int, tab_id: int, window_id: int, path: str, position: str, alpha: float) -> None: +def set_window_logo(os_window_id: int, tab_id: int, window_id: int, path: str, position: str, alpha: float, png_data: bytes = b'') -> None: pass diff --git a/kitty/rc/base.py b/kitty/rc/base.py index e3041b889..3fcc11a02 100644 --- a/kitty/rc/base.py +++ b/kitty/rc/base.py @@ -1,9 +1,9 @@ #!/usr/bin/env python # License: GPLv3 Copyright: 2020, Kovid Goyal -import tempfile from contextlib import suppress from dataclasses import dataclass, field +from io import BytesIO from typing import TYPE_CHECKING, Any, Callable, Dict, FrozenSet, Iterable, Iterator, List, NoReturn, Optional, Set, Tuple, Type, Union, cast from kitty.cli import CompletionSpec, get_defaults_from_seq, parse_args, parse_option_spec @@ -27,16 +27,6 @@ class NoResponse: pass -class NamedTemporaryFile: - name: str = '' - - def __enter__(self) -> None: ... - def __exit__(self, exc: Any, value: Any, tb: Any) -> None: ... - def close(self) -> None: ... - def write(self, data: bytes) -> None: ... - def flush(self) -> None: ... - - class RemoteControlError(Exception): pass @@ -279,22 +269,27 @@ class StreamInFlight: def __init__(self) -> None: self.stream_id = '' - self.tempfile: Optional[NamedTemporaryFile] = None + self.tempfile: Optional[BytesIO] = None - def handle_data(self, stream_id: str, data: bytes) -> Union[AsyncResponse, NamedTemporaryFile]: + def handle_data(self, stream_id: str, data: bytes) -> Union[AsyncResponse, BytesIO]: from ..remote_control import close_active_stream - if stream_id != self.stream_id: + def abort_stream() -> None: close_active_stream(self.stream_id) + self.stream_id = '' if self.tempfile is not None: self.tempfile.close() self.tempfile = None + + if stream_id != self.stream_id: + abort_stream() self.stream_id = stream_id if self.tempfile is None: - t: NamedTemporaryFile = cast(NamedTemporaryFile, tempfile.NamedTemporaryFile(suffix='.png')) - self.tempfile = t - else: - t = self.tempfile + self.tempfile = BytesIO() + t = self.tempfile if data: + if (t.tell() + len(data)) > 128 * 1024 * 1024: + abort_stream() + raise StreamError('Too much data being sent') t.write(data) return AsyncResponse() close_active_stream(self.stream_id) @@ -404,7 +399,7 @@ class RemoteCommand: def cancel_async_request(self, boss: 'Boss', window: Optional['Window'], payload_get: PayloadGetType) -> None: pass - def handle_streamed_data(self, data: bytes, payload_get: PayloadGetType) -> Union[NamedTemporaryFile, AsyncResponse]: + def handle_streamed_data(self, data: bytes, payload_get: PayloadGetType) -> Union[BytesIO, AsyncResponse]: stream_id = payload_get('stream_id') if not stream_id or not isinstance(stream_id, str): raise StreamError('No stream_id in rc payload') diff --git a/kitty/rc/set_background_image.py b/kitty/rc/set_background_image.py index 40560259f..36cb70eeb 100644 --- a/kitty/rc/set_background_image.py +++ b/kitty/rc/set_background_image.py @@ -3,6 +3,7 @@ import os from base64 import standard_b64decode, standard_b64encode +from io import BytesIO from typing import TYPE_CHECKING, Optional from kitty.types import AsyncResponse @@ -14,7 +15,6 @@ from .base import ( Boss, CmdGenerator, ImageCompletion, - NamedTemporaryFile, PayloadGetType, PayloadType, RCOptions, @@ -111,17 +111,16 @@ failed, the command will exit with a success code. layout = payload_get('layout') if data == '-': path = None - tfile = NamedTemporaryFile() + tfile = BytesIO() else: q = self.handle_streamed_data(standard_b64decode(data) if data else b'', payload_get) if isinstance(q, AsyncResponse): return q - path = q.name + path = '/image/from/remote/control' tfile = q try: - with tfile: - boss.set_background_image(path, os_windows, payload_get('configured'), layout) + boss.set_background_image(path, os_windows, payload_get('configured'), layout, tfile.getvalue()) except ValueError as err: err.hide_traceback = True # type: ignore raise diff --git a/kitty/rc/set_window_logo.py b/kitty/rc/set_window_logo.py index 73bb3f811..ba20ee3eb 100644 --- a/kitty/rc/set_window_logo.py +++ b/kitty/rc/set_window_logo.py @@ -3,6 +3,7 @@ import os from base64 import standard_b64decode, standard_b64encode +from io import BytesIO from typing import TYPE_CHECKING, Optional from kitty.types import AsyncResponse @@ -14,7 +15,6 @@ from .base import ( Boss, CmdGenerator, ImageCompletion, - NamedTemporaryFile, PayloadGetType, PayloadType, RCOptions, @@ -105,18 +105,18 @@ failed, the command will exit with a success code. position = payload_get('position') or '' if data == '-': path = '' - tfile = NamedTemporaryFile() + tfile = BytesIO() else: q = self.handle_streamed_data(standard_b64decode(data) if data else b'', payload_get) if isinstance(q, AsyncResponse): return q - path = q.name + import hashlib + path = '/from/remote/control/' + hashlib.sha1(q.getvalue()).hexdigest() tfile = q - with tfile: - for window in self.windows_for_match_payload(boss, window, payload_get): - if window: - window.set_logo(path, position, alpha) + for window in self.windows_for_match_payload(boss, window, payload_get): + if window: + window.set_logo(path, position, alpha, tfile.getvalue()) return None diff --git a/kitty/state.c b/kitty/state.c index bb5f18fc3..8b69cfb1f 100644 --- a/kitty/state.c +++ b/kitty/state.c @@ -256,10 +256,10 @@ release_gpu_resources_for_window(Window *w) { } static bool -set_window_logo(Window *w, const char *path, const ImageAnchorPosition pos, float alpha, bool is_default) { +set_window_logo(Window *w, const char *path, const ImageAnchorPosition pos, float alpha, bool is_default, char *png_data, size_t png_data_size) { bool ok = false; if (path && path[0]) { - window_logo_id_t wl = find_or_create_window_logo(global_state.all_window_logos, path); + window_logo_id_t wl = find_or_create_window_logo(global_state.all_window_logos, path, png_data, png_data_size); if (wl) { if (w->window_logo.id) decref_window_logo(global_state.all_window_logos, w->window_logo.id); w->window_logo.id = wl; @@ -285,7 +285,7 @@ initialize_window(Window *w, PyObject *title, bool init_gpu_resources) { w->visible = true; w->title = title; Py_XINCREF(title); - if (!set_window_logo(w, OPT(default_window_logo), OPT(window_logo_position), OPT(window_logo_alpha), true)) { + if (!set_window_logo(w, OPT(default_window_logo), OPT(window_logo_position), OPT(window_logo_alpha), true, NULL, 0)) { log_error("Failed to load default window logo: %s", OPT(default_window_logo)); if (PyErr_Occurred()) PyErr_Print(); } @@ -1106,7 +1106,7 @@ PYWRAP0(apply_options_update) { for (size_t w = 0; w < tab->num_windows; w++) { Window *window = tab->windows + w; if (window->window_logo.using_default) { - set_window_logo(window, OPT(default_window_logo), OPT(window_logo_position), OPT(window_logo_alpha), true); + set_window_logo(window, OPT(default_window_logo), OPT(window_logo_position), OPT(window_logo_alpha), true, NULL, 0); } } } @@ -1153,7 +1153,7 @@ pyset_background_image(PyObject *self UNUSED, PyObject *args) { PyObject *os_window_ids; int configured = 0; char *png_data = NULL; Py_ssize_t png_data_size = 0; - PA("zO!|pOy#", &path, &PyTuple_Type, &os_window_ids, &configured, &layout_name, &png_data, png_data_size); + PA("zO!|pOy#", &path, &PyTuple_Type, &os_window_ids, &configured, &layout_name, &png_data, &png_data_size); size_t size; BackgroundImageLayout layout = PyUnicode_Check(layout_name) ? bglayout(layout_name) : OPT(background_image_layout); BackgroundImage *bgimage = NULL; @@ -1169,6 +1169,7 @@ pyset_background_image(PyObject *self UNUSED, PyObject *args) { return NULL; } ok = png_from_file_pointer(fp, path, &bgimage->bitmap, &bgimage->width, &bgimage->height, &size); + fclose(fp); } else { ok = png_path_to_bitmap(path, &bgimage->bitmap, &bgimage->width, &bgimage->height, &size); } @@ -1280,10 +1281,11 @@ PYWRAP1(set_window_logo) { id_type os_window_id, tab_id, window_id; const char *path; PyObject *position; float alpha = 0.5; - PA("KKKsUf", &os_window_id, &tab_id, &window_id, &path, &position, &alpha); + char *png_data = NULL; Py_ssize_t png_data_size = 0; + PA("KKKsUf|y#", &os_window_id, &tab_id, &window_id, &path, &position, &alpha, &png_data, &png_data_size); bool ok = false; WITH_WINDOW(os_window_id, tab_id, window_id); - ok = set_window_logo(window, path, PyObject_IsTrue(position) ? bganchor(position) : OPT(window_logo_position), (0 <= alpha && alpha <= 1) ? alpha : OPT(window_logo_alpha), false); + ok = set_window_logo(window, path, PyObject_IsTrue(position) ? bganchor(position) : OPT(window_logo_position), (0 <= alpha && alpha <= 1) ? alpha : OPT(window_logo_alpha), false, png_data, png_data_size); END_WITH_WINDOW; if (ok) Py_RETURN_TRUE; Py_RETURN_FALSE; diff --git a/kitty/window.py b/kitty/window.py index d8ee47ed4..0fd85a42e 100644 --- a/kitty/window.py +++ b/kitty/window.py @@ -1449,9 +1449,9 @@ class Window: 'text': text } - def set_logo(self, path: str, position: str = '', alpha: float = -1) -> None: + def set_logo(self, path: str, position: str = '', alpha: float = -1, png_data: bytes = b'') -> None: path = resolve_custom_file(path) if path else '' - set_window_logo(self.os_window_id, self.tab_id, self.id, path, position or '', alpha) + set_window_logo(self.os_window_id, self.tab_id, self.id, path, position or '', alpha, png_data) def paste_with_actions(self, text: str) -> None: if self.destroyed or not text: diff --git a/kitty/window_logo.c b/kitty/window_logo.c index 825f7e90b..2beff81a1 100644 --- a/kitty/window_logo.c +++ b/kitty/window_logo.c @@ -49,7 +49,7 @@ set_on_gpu_state(WindowLogo *s, bool on_gpu) { } window_logo_id_t -find_or_create_window_logo(WindowLogoTable *head, const char *path) { +find_or_create_window_logo(WindowLogoTable *head, const char *path, void *png_data, size_t png_data_size) { WindowLogoItem *s = NULL; unsigned _uthash_hfstr_keylen = (unsigned)uthash_strlen(path); HASH_FIND(hh_path, head->by_path, path, _uthash_hfstr_keylen, s); @@ -62,7 +62,17 @@ find_or_create_window_logo(WindowLogoTable *head, const char *path) { if (!s) { PyErr_NoMemory(); return 0; } s->path = strdup(path); if (!s->path) { free(s); PyErr_NoMemory(); return 0; } - if (png_path_to_bitmap(path, &s->wl.bitmap, &s->wl.width, &s->wl.height, &size)) s->wl.load_from_disk_ok = true; + bool ok = false; + if (png_data == NULL) { + ok = png_path_to_bitmap(path, &s->wl.bitmap, &s->wl.width, &s->wl.height, &size); + } else { + FILE *fp = fmemopen(png_data, png_data_size, "r"); + if (fp != NULL) { + ok = png_from_file_pointer(fp, path, &s->wl.bitmap, &s->wl.width, &s->wl.height, &size); + fclose(fp); + } + } + if (ok) s->wl.load_from_disk_ok = true; s->refcnt++; static window_logo_id_t idc = 0; s->id = ++idc; diff --git a/kitty/window_logo.h b/kitty/window_logo.h index 711b46e24..d349a4400 100644 --- a/kitty/window_logo.h +++ b/kitty/window_logo.h @@ -20,7 +20,7 @@ typedef struct WindowLogo { typedef struct WindowLogoTable WindowLogoTable; window_logo_id_t -find_or_create_window_logo(WindowLogoTable *table, const char *path); +find_or_create_window_logo(WindowLogoTable *table, const char *path, void *png_data, size_t png_data_size); WindowLogo* find_window_logo(WindowLogoTable *table, window_logo_id_t id);