Add *_into API for streaming encode/decode

kitty/data-types.c

@@ -201,6 +201,21 @@ StreamingBase64Encoder_encode(StreamingBase64Decoder *self, PyObject *a) {
     return Py_NewRef(ans);
 }
 
+static PyObject*
+StreamingBase64Encoder_encode_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_encode(src.len);
+    if ((Py_ssize_t)sz > data.len) { PyErr_SetString(PyExc_BufferError, "output buffer too small"); return NULL; }
+    base64_stream_encode(&self->state, src.buf, src.len, data.buf, &sz);
+    return PyLong_FromSize_t(sz);
+}
+
 static PyObject*
 StreamingBase64Encoder_reset(StreamingBase64Decoder *self, PyObject *args UNUSED) {
     char trailer[4];
@@ -219,6 +234,7 @@ static PyTypeObject StreamingBase64Encoder_Type = {
     .tp_doc = "StreamingBase64Encoder",
     .tp_methods = (PyMethodDef[]){
         {"encode", (PyCFunction)StreamingBase64Encoder_encode, METH_O, ""},
+        {"encode_into", (PyCFunction)(void(*)(void))StreamingBase64Encoder_encode_into, METH_FASTCALL, ""},
         {"reset", (PyCFunction)StreamingBase64Encoder_reset, METH_NOARGS, ""},
         {NULL, NULL, 0, NULL},
     },

kitty/fast_data_types.pyi

@@ -8,7 +8,7 @@ from kitty.fonts.render import FontObject
 from kitty.marks import MarkerFunc
 from kitty.options.types import Options
 from kitty.types import LayerShellConfig, SignalInfo
-from kitty.typing import EdgeLiteral, NotRequired
+from kitty.typing import EdgeLiteral, NotRequired, ReadableBuffer, WriteableBuffer
 
 # Constants {{{
 GLFW_LAYER_SHELL_NONE: int
@@ -305,9 +305,6 @@ WINDOW_MINIMIZED: int
 # }}}
 
 
-ReadOnlyBuffer = Union[bytes, bytearray, memoryview]
-
-
 def encode_key_for_tty(
     key: int = 0,
     shifted_key: int = 0,
@@ -1705,18 +1702,23 @@ def get_mouse_data_for_window(os_window_id: int, tab_id: int, window_id: int) ->
 
 
 class StreamingBase64Decoder:
-    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
+    # reset the state to empty to start decoding a new stream
+    def reset(self) -> None: ...
+    # decode the specified data
+    def decode(self, data: ReadableBuffer) -> bytes: ...
+    # decode the specified data, return number of bytes written dest should be as large as src (technically 3/4 src + 2)
+    def decode_into(self, dest: WriteableBuffer, src: ReadableBuffer) -> int: ...
 
 
 
 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
+    # encode the specified data
+    def encode(self, data: ReadableBuffer) -> bytes: ...
+    # reset the state to empty to start encoding a new stream, return any trailing bytes from the previous encode call
+    def reset(self) -> bytes: ...
+    # encode the specified data, return number of bytes written dest should be at least 4/3 *src + 2 bytes in size
+    def encode_into(self, dest: WriteableBuffer, src: ReadableBuffer) -> int: ...
 
 
 

kitty/file_transmission.py

@@ -22,6 +22,7 @@ from typing import IO, Any, Callable, DefaultDict, Deque, Dict, Iterable, Iterat
 from kittens.transfer.utils import IdentityCompressor, ZlibCompressor, abspath, expand_home, home_path
 from kitty.fast_data_types import ESC_OSC, FILE_TRANSFER_CODE, AES256GCMDecrypt, add_timer, base64_decode, base64_encode, get_boss, get_options, monotonic
 from kitty.types import run_once
+from kitty.typing import ReadableBuffer, WriteableBuffer
 
 from .utils import log_error
 
@@ -412,12 +413,12 @@ class PatchFile:
             return os.path.getsize(self.path)
         return df.tell()
 
-    def read_from_src(self, pos: int, b: Union[bytearray, memoryview]) -> int:
+    def read_from_src(self, pos: int, b: WriteableBuffer) -> int:
         assert self.src_file is not None
         self.src_file.seek(pos, os.SEEK_SET)
         return self.src_file.readinto(b)
 
-    def write_to_dest(self, b: Union[bytes, bytearray, memoryview]) -> None:
+    def write_to_dest(self, b: ReadableBuffer) -> None:
         self.dest_file.write(b)
 
     def write(self, b: bytes) -> None:
@@ -669,7 +670,7 @@ class SourceFile:
         self.buf = bytearray()
         self.write_pos = 0
 
-    def write(self, b: Union[bytes, bytearray, memoryview]) -> None:
+    def write(self, b: ReadableBuffer) -> None:
         self.buf[self.write_pos:self.write_pos+len(b)] = b
         self.write_pos += len(b)
 

kitty/typing.py

@@ -12,6 +12,7 @@ KeyEventType = ImageManagerType = KittyCommonOpts = HandlerType = None
 GRT_t = GRT_a = GRT_d = GRT_f = GRT_m = GRT_o = GRT_C = None
 ScreenSize = KittensKeyActionType = MouseEvent = MouseButton = AbstractEventLoop = None
 TermManagerType = LoopType = Debug = GraphicsCommandType = None
+ReadableBuffer = WriteableBuffer = bytearray
 
 CompletedProcess = Tuple
 TypedDict = dict

kitty/typing.pyi

@@ -1,3 +1,5 @@
+import array
+import mmap
 from asyncio import AbstractEventLoop as AbstractEventLoop
 from socket import AddressFamily as AddressFamily
 from socket import socket as Socket
@@ -50,6 +52,9 @@ GRT_o = Literal['z', 'z']  # two z's to workaround a bug in ruff
 GRT_m = Literal[0, 1]
 GRT_C = Literal[0, 1]
 GRT_d = Literal['a', 'A', 'c', 'C', 'i', 'I', 'p', 'P', 'q', 'Q', 'x', 'X', 'y', 'Y', 'z', 'Z', 'f', 'F']
+ReadableBuffer = bytes | bytearray | memoryview | array.array[int] | mmap.mmap
+WriteableBuffer = bytearray | memoryview | array.array[int] | mmap.mmap
+
 
 
 class WindowSystemMouseEvent(TypedDict):
@@ -65,5 +70,6 @@ __all__ = (
     'TermManagerType', 'BossType', 'ChildType', 'BadLineType', 'MouseButton', 'NotRequired',
     'KeyActionType', 'KeyMap', 'KittyCommonOpts', 'AliasMap', 'CoreTextFont', 'WindowSystemMouseEvent',
     'FontConfigPattern', 'ScreenType', 'StartupCtx', 'KeyEventType', 'LayoutType', 'PowerlineStyle',
-    'RemoteCommandType', 'SessionType', 'SessionTab', 'SpecialWindowInstance', 'TabType', 'ScreenSize', 'WindowType'
+    'RemoteCommandType', 'SessionType', 'SessionTab', 'SpecialWindowInstance', 'TabType', 'ScreenSize', 'WindowType',
+    'ReadableBuffer', 'WriteableBuffer',
 )