From d2ced8d6e148f28d319aa46f9300a9a5f8f491fb Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 29 Jul 2024 22:08:13 +0530
Subject: [PATCH] Add a streaming base64 encoder

---
 kitty/data-types.c        | 55 +++++++++++++++++++++++++++++++++++++++
 kitty/fast_data_types.pyi |  7 +++++
 2 files changed, 62 insertions(+)

diff --git a/kitty/data-types.c b/kitty/data-types.c
index 05117826c..61c80bbbc 100644
--- a/kitty/data-types.c
+++ b/kitty/data-types.c
@@ -108,6 +108,7 @@ pybase64_decode(PyObject UNUSED *self, PyObject *args) {
 typedef struct StreamingBase64Decoder {
     PyObject_HEAD
     struct base64_state state;
+    bool add_trailing_bytes;
 } StreamingBase64Decoder;
 
 static int
@@ -154,6 +155,58 @@ static PyTypeObject StreamingBase64Decoder_Type = {
     .tp_init = StreamingBase64Decoder_init,
 };
 
+static int
+StreamingBase64Encoder_init(PyObject *s, PyObject *args, PyObject *kwds UNUSED) {
+    StreamingBase64Decoder *self = (StreamingBase64Decoder*)s;
+    self->add_trailing_bytes = true;
+    switch (PyTuple_GET_SIZE(args)) {
+        case 0: break;
+        case 1: self->add_trailing_bytes = PyObject_IsTrue(PyTuple_GET_ITEM(args, 0)); break;
+        default: PyErr_SetString(PyExc_TypeError, "constructor takes no more than one argument"); return -1;
+    }
+	base64_stream_encode_init(&self->state, 0);
+    return 0;
+}
+
+static PyObject*
+StreamingBase64Encoder_encode(StreamingBase64Decoder *self, PyObject *a) {
+    RAII_PY_BUFFER(data);
+    if (PyObject_GetBuffer(a, &data, PyBUF_SIMPLE) != 0) return NULL;
+    if (!data.buf || !data.len) return PyBytes_FromStringAndSize(NULL, 0);
+    size_t sz = required_buffer_size_for_base64_encode(data.len);
+    RAII_PyObject(ans, PyBytes_FromStringAndSize(NULL, sz));
+    if (!ans) return NULL;
+    base64_stream_encode(&self->state, data.buf, data.len, PyBytes_AS_STRING(ans), &sz);
+    if (_PyBytes_Resize(&ans, sz) != 0) return NULL;
+    return Py_NewRef(ans);
+}
+
+static PyObject*
+StreamingBase64Encoder_reset(StreamingBase64Decoder *self, PyObject *args UNUSED) {
+    char trailer[4];
+    size_t sz;
+    base64_stream_encode_final(&self->state, trailer, &sz);
+    base64_stream_encode_init(&self->state, 0);
+    if (!self->add_trailing_bytes) { while(sz && trailer[sz-1] == '=') sz--; }
+    return PyBytes_FromStringAndSize(trailer, sz);
+}
+
+static PyTypeObject StreamingBase64Encoder_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    .tp_name = "kitty.fast_data_types.StreamingBase64Encoder",
+    .tp_basicsize = sizeof(StreamingBase64Decoder),
+    .tp_flags = Py_TPFLAGS_DEFAULT,
+    .tp_doc = "StreamingBase64Encoder",
+    .tp_methods = (PyMethodDef[]){
+        {"encode", (PyCFunction)StreamingBase64Encoder_encode, METH_O, ""},
+        {"reset", (PyCFunction)StreamingBase64Encoder_reset, METH_NOARGS, ""},
+        {NULL, NULL, 0, NULL},
+    },
+    .tp_new = PyType_GenericNew,
+    .tp_init = StreamingBase64Encoder_init,
+};
+
+
 static PyObject*
 pyset_iutf8(PyObject UNUSED *self, PyObject *args) {
     int fd, on;
@@ -661,6 +714,8 @@ PyInit_fast_data_types(void) {
 
     if (PyType_Ready(&StreamingBase64Decoder_Type) < 0) return NULL;
     if (PyModule_AddObject(m, "StreamingBase64Decoder", (PyObject *) &StreamingBase64Decoder_Type) < 0) return NULL;
+    if (PyType_Ready(&StreamingBase64Encoder_Type) < 0) return NULL;
+    if (PyModule_AddObject(m, "StreamingBase64Encoder", (PyObject *) &StreamingBase64Encoder_Type) < 0) return NULL;
 
     return m;
 }
diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi
index 30492d5e3..534ab95c7 100644
--- a/kitty/fast_data_types.pyi
+++ b/kitty/fast_data_types.pyi
@@ -1709,6 +1709,13 @@ class StreamingBase64Decoder:
     def reset(self) -> None: ...  # reset the state to empty to start decoding a new stream
 
 
+class StreamingBase64Encodeer:
+    def __init__(self, add_trailing_bytes: bool = True) -> None: ...
+    def encode(self, data: ReadOnlyBuffer) -> bytes: ...  # decode the specified data
+    def reset(self) -> bytes: ...  # reset the state to empty to start decoding a new stream, return any trailing bytes
+
+
+
 class DiskCache:
     small_hole_threshold: int
     defrag_factor: int