From b52275e0b527c7671ab5b8449fbb17b929af2712 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Mon, 29 Jul 2024 21:24:45 +0530 Subject: [PATCH] Simplify API of streaming base64 decoder further --- kitty/clipboard.py | 12 +++--- kitty/data-types.c | 81 ++++++++------------------------------- kitty/fast_data_types.pyi | 8 +--- kitty/notifications.py | 10 ++--- 4 files changed, 26 insertions(+), 85 deletions(-) diff --git a/kitty/clipboard.py b/kitty/clipboard.py index 2cf52ffb2..e74dfa0cb 100644 --- a/kitty/clipboard.py +++ b/kitty/clipboard.py @@ -237,7 +237,7 @@ class WriteRequest: self, is_primary_selection: bool = False, protocol_type: ProtocolType = ProtocolType.osc_52, id: str = '', rollover_size: int = 16 * 1024 * 1024, max_size: int = -1, ) -> None: - self.decoder = StreamingBase64Decoder(8 * 1024) + self.decoder = StreamingBase64Decoder() self.id = id self.is_primary_selection = is_primary_selection self.protocol_type = protocol_type @@ -281,18 +281,16 @@ class WriteRequest: def flush_base64_data(self) -> None: if self.currently_writing_mime: - self.decoder.reinitialize() - if len(self.decoder): - self.write_base64_data(b'') + self.decoder.reset() start = self.mime_map[self.currently_writing_mime][0] self.mime_map[self.currently_writing_mime] = MimePos(start, self.tempfile.tell() - start) self.currently_writing_mime = '' def write_base64_data(self, b: bytes) -> None: if not self.max_size_exceeded: - self.decoder.add(b) - if len(self.decoder): - self.tempfile.write(self.decoder.take_output()) + decoded = self.decoder.decode(b) + if decoded: + self.tempfile.write(decoded) if self.max_size > 0 and self.tempfile.tell() > (self.max_size * 1024 * 1024): log_error(f'Clipboard write request has more data than allowed by clipboard_max_size ({self.max_size}), truncating') self.max_size_exceeded = True diff --git a/kitty/data-types.c b/kitty/data-types.c index 16347d6e4..a53eb5973 100644 --- a/kitty/data-types.c +++ b/kitty/data-types.c @@ -107,100 +107,51 @@ pybase64_decode(PyObject UNUSED *self, PyObject *args) { typedef struct StreamingBase64Decoder { PyObject_HEAD - PyObject *output; - size_t output_sz, output_capacity, initial_capacity; struct base64_state state; } StreamingBase64Decoder; static int -StreamingBase64Decoder_init(PyObject *s, PyObject *args, PyObject *kwds) { - static char *kwlist[] = {"initial_capacity", NULL}; - unsigned long initial_capacity = 8 * 1024; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|k", kwlist, &initial_capacity)) return -1; +StreamingBase64Decoder_init(PyObject *s, PyObject *args UNUSED, PyObject *kwds UNUSED) { + if (PyTuple_GET_SIZE(args)) { PyErr_SetString(PyExc_TypeError, "constructor takes no arguments"); return -1; } StreamingBase64Decoder *self = (StreamingBase64Decoder*)s; - self->initial_capacity = initial_capacity; base64_stream_decode_init(&self->state, 0); return 0; } -static void -StreamingBase64Decoder_dealloc(PyObject *self) { - StreamingBase64Decoder *h = (StreamingBase64Decoder*)self; - Py_CLEAR(h->output); - Py_TYPE(self)->tp_free(self); -} - -static bool -write_base64_data(StreamingBase64Decoder *self, const void *data, const size_t len) { - if (!len) return true; - size_t sz = required_buffer_size_for_base64_decode(len); - if ((self->output_sz + sz) > self->output_capacity) { - size_t cap = MAX(self->output_capacity * 2, MAX(self->output_sz + sz, self->initial_capacity)); - if (self->output) { if (_PyBytes_Resize(&self->output, cap) != 0) return false; } - else { self->output = PyBytes_FromStringAndSize(NULL, cap); if (!self->output) return false; } - self->output_capacity = cap; - } - bool ok = base64_stream_decode(&self->state, data, len, PyBytes_AS_STRING(self->output) + self->output_sz, &sz); - if (!ok) { - PyErr_SetString(PyExc_ValueError, "Invalid base64 input data"); - return false; - } - self->output_sz += sz; - return true; -} - static PyObject* -StreamingBase64Decoder_add(StreamingBase64Decoder *self, PyObject *a) { +StreamingBase64Decoder_decode(StreamingBase64Decoder *self, PyObject *a) { RAII_PY_BUFFER(data); if (PyObject_GetBuffer(a, &data, PyBUF_SIMPLE) != 0) return NULL; - if (!data.buf || !data.len) return PyLong_FromLong(0); - size_t before = self->output_sz; - if (!write_base64_data(self, data.buf, data.len)) return NULL; - return PyLong_FromSize_t(self->output_sz - before); + if (!data.buf || !data.len) return PyBytes_FromStringAndSize(NULL, 0); + size_t sz = required_buffer_size_for_base64_decode(data.len); + RAII_PyObject(ans, PyBytes_FromStringAndSize(NULL, sz)); + if (!ans) return NULL; + if (!base64_stream_decode(&self->state, data.buf, data.len, PyBytes_AS_STRING(ans), &sz)) { + PyErr_SetString(PyExc_ValueError, "Invalid base64 input data"); + return NULL; + } + if (_PyBytes_Resize(&ans, sz) != 0) return NULL; + return Py_NewRef(ans); } -static Py_ssize_t -StreamingBase64Decoder_len(PyObject *s) { return ((StreamingBase64Decoder*)s)->output_sz; } - static PyObject* -StreamingBase64Decoder_reinitialize(StreamingBase64Decoder *self, PyObject *args UNUSED) { +StreamingBase64Decoder_reset(StreamingBase64Decoder *self, PyObject *args UNUSED) { base64_stream_decode_init(&self->state, 0); Py_RETURN_NONE; } - -static PyObject* -StreamingBase64Decoder_copy_output(StreamingBase64Decoder *self, PyObject *args UNUSED) { - return PyBytes_FromStringAndSize(PyBytes_AS_STRING(self->output), self->output_sz); -} - -static PyObject* -StreamingBase64Decoder_take_output(StreamingBase64Decoder *self, PyObject *args UNUSED) { - if (!self->output_sz) return PyBytes_FromStringAndSize(NULL, 0); - if (_PyBytes_Resize(&self->output, self->output_sz) != 0) return NULL; - PyObject *ans = self->output; - self->output = NULL; self->output_sz = 0; self->output_capacity = 0; - return ans; -} - static PyTypeObject StreamingBase64Decoder_Type = { PyVarObject_HEAD_INIT(NULL, 0) .tp_name = "kitty.fast_data_types.StreamingBase64Decoder", .tp_basicsize = sizeof(StreamingBase64Decoder), - .tp_dealloc = StreamingBase64Decoder_dealloc, .tp_flags = Py_TPFLAGS_DEFAULT, .tp_doc = "StreamingBase64Decoder", .tp_methods = (PyMethodDef[]){ - {"add", (PyCFunction)StreamingBase64Decoder_add, METH_O, ""}, - {"reinitialize", (PyCFunction)StreamingBase64Decoder_reinitialize, METH_NOARGS, ""}, - {"take_output", (PyCFunction)StreamingBase64Decoder_take_output, METH_NOARGS, ""}, - {"copy_output", (PyCFunction)StreamingBase64Decoder_copy_output, METH_NOARGS, ""}, + {"decode", (PyCFunction)StreamingBase64Decoder_decode, METH_O, ""}, + {"reset", (PyCFunction)StreamingBase64Decoder_reset, METH_NOARGS, ""}, {NULL, NULL, 0, NULL}, }, .tp_new = PyType_GenericNew, .tp_init = StreamingBase64Decoder_init, - .tp_as_sequence = &(PySequenceMethods){ - .sq_length = StreamingBase64Decoder_len, - }, }; static PyObject* diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi index 33747fe40..30492d5e3 100644 --- a/kitty/fast_data_types.pyi +++ b/kitty/fast_data_types.pyi @@ -1705,12 +1705,8 @@ def get_mouse_data_for_window(os_window_id: int, tab_id: int, window_id: int) -> class StreamingBase64Decoder: - def __init__(self, initial_capacity: int = 8 *1024) -> None: ... # set the initial output buffer capacity - def add(self, data: ReadOnlyBuffer) -> int: ... # add the base64 data - def reinitialize(self) -> None: ... # reset the state to empty to start decoding a new stream - def take_output(self) -> bytes: ... # take the output so far. The decoder no longer references this output - def copy_output(self) -> bytes: ... # copy the output so far - def __len__(self) -> int: ... # return the length of the current output + def decode(self, data: ReadOnlyBuffer) -> bytes: ... # decode the specified data + def reset(self) -> None: ... # reset the state to empty to start decoding a new stream class DiskCache: diff --git a/kitty/notifications.py b/kitty/notifications.py index a09f44a71..987c797f9 100644 --- a/kitty/notifications.py +++ b/kitty/notifications.py @@ -162,7 +162,7 @@ class DataStore: class EncodedDataStore: def __init__(self, data_store: DataStore) -> None: - self.decoder = StreamingBase64Decoder(initial_capacity=4096) + self.decoder = StreamingBase64Decoder() self.data_store = data_store @property @@ -178,14 +178,10 @@ class EncodedDataStore: def add_base64_data(self, data: Union[str, bytes]) -> None: if isinstance(data, str): data = data.encode('ascii') - self.decoder.add(data) - if len(self.decoder) >= self.data_store.max_size: - self.data_store(self.decoder.take_output()) + self.data_store(self.decoder.decode(data)) def flush_encoded_data(self) -> None: - self.decoder.reinitialize() - if len(self.decoder): - self.data_store(self.decoder.take_output()) + self.decoder.reset() def finalise(self) -> bytes: self.flush_encoded_data()