Add a decode_into API for streaming base64 decoding

Allows buffer reuse
This commit is contained in:
Kovid Goyal
2024-07-29 22:30:36 +05:30
parent d2ced8d6e1
commit 8d36677d41
2 changed files with 25 additions and 1 deletions

View File

@@ -135,11 +135,30 @@ StreamingBase64Decoder_decode(StreamingBase64Decoder *self, PyObject *a) {
return Py_NewRef(ans);
}
static PyObject*
StreamingBase64Decoder_decode_into(StreamingBase64Decoder *self, PyObject *const *args, Py_ssize_t nargs) {
if (nargs != 2) { PyErr_SetString(PyExc_TypeError, "constructor takes exactly two arguments"); return NULL; }
RAII_PY_BUFFER(data);
if (PyObject_GetBuffer(args[0], &data, PyBUF_WRITE) != 0) return NULL;
if (!data.buf || !data.len) return PyLong_FromLong(0);
RAII_PY_BUFFER(src);
if (PyObject_GetBuffer(args[1], &src, PyBUF_SIMPLE) != 0) return NULL;
if (!src.buf || !src.len) return PyLong_FromLong(0);
size_t sz = required_buffer_size_for_base64_decode(src.len);
if ((Py_ssize_t)sz > data.len) { PyErr_SetString(PyExc_BufferError, "output buffer too small"); return NULL; }
if (!base64_stream_decode(&self->state, src.buf, src.len, data.buf, &sz)) {
PyErr_SetString(PyExc_ValueError, "Invalid base64 input data");
return NULL;
}
return PyLong_FromSize_t(sz);
}
static PyObject*
StreamingBase64Decoder_reset(StreamingBase64Decoder *self, PyObject *args UNUSED) {
base64_stream_decode_init(&self->state, 0);
Py_RETURN_NONE;
}
static PyTypeObject StreamingBase64Decoder_Type = {
PyVarObject_HEAD_INIT(NULL, 0)
.tp_name = "kitty.fast_data_types.StreamingBase64Decoder",
@@ -148,6 +167,7 @@ static PyTypeObject StreamingBase64Decoder_Type = {
.tp_doc = "StreamingBase64Decoder",
.tp_methods = (PyMethodDef[]){
{"decode", (PyCFunction)StreamingBase64Decoder_decode, METH_O, ""},
{"decode_into", (PyCFunction)(void(*)(void))StreamingBase64Decoder_decode_into, METH_FASTCALL, ""},
{"reset", (PyCFunction)StreamingBase64Decoder_reset, METH_NOARGS, ""},
{NULL, NULL, 0, NULL},
},

View File

@@ -1705,8 +1705,12 @@ def get_mouse_data_for_window(os_window_id: int, tab_id: int, window_id: int) ->
class StreamingBase64Decoder:
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
def decode(self, data: ReadOnlyBuffer) -> bytes: ... # decode the specified data
def decode_into(self, dest: bytearray, src: ReadOnlyBuffer) -> int: ... # decode the specified data, return number of bytes written
# dest should be as large as src (technically 3/4 src + 2)
# dest can be any writable buffer
class StreamingBase64Encodeer: