From e97c225b6c2afccf8a15999b50fc6c6d0cb44974 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Tue, 30 Jul 2024 07:18:33 +0530 Subject: [PATCH] Add *_into API for streaming encode/decode --- kittens/transfer/rsync.pyi | 23 +++++++++++------------ kitty/data-types.c | 16 ++++++++++++++++ kitty/fast_data_types.pyi | 24 +++++++++++++----------- kitty/file_transmission.py | 7 ++++--- kitty/typing.py | 1 + kitty/typing.pyi | 8 +++++++- 6 files changed, 52 insertions(+), 27 deletions(-) diff --git a/kittens/transfer/rsync.pyi b/kittens/transfer/rsync.pyi index 545553b78..8e79dea4c 100644 --- a/kittens/transfer/rsync.pyi +++ b/kittens/transfer/rsync.pyi @@ -1,14 +1,13 @@ from typing import Callable, Union -ReadOnlyBuffer = Union[bytes, bytearray, memoryview] -WriteBuffer = Union[bytearray, memoryview] +from kitty.typing import ReadableBuffer, WriteableBuffer class RsyncError(Exception): pass class Hasher: - def __init__(self, which: str, data: ReadOnlyBuffer = b''): ... - def update(self, data: ReadOnlyBuffer) -> None: ... + def __init__(self, which: str, data: ReadableBuffer = b''): ... + def update(self, data: ReadableBuffer) -> None: ... def reset(self) -> None: ... def digest(self) -> bytes: ... def hexdigest(self) -> str: ... @@ -20,16 +19,16 @@ class Hasher: @property def name(self) -> str: ... -def xxh128_hash(data: ReadOnlyBuffer) -> bytes: ... -def xxh128_hash_with_seed(data: ReadOnlyBuffer, seed: int) -> bytes: ... +def xxh128_hash(data: ReadableBuffer) -> bytes: ... +def xxh128_hash_with_seed(data: ReadableBuffer, seed: int) -> bytes: ... class Patcher: def __init__(self, expected_input_size: int = 0): ... - def signature_header(self, output: WriteBuffer) -> int: ... - def sign_block(self, block: ReadOnlyBuffer, output: WriteBuffer) -> int: ... - def apply_delta_data(self, data: ReadOnlyBuffer, read: Callable[[int, WriteBuffer], int], write: Callable[[ReadOnlyBuffer], None]) -> None: ... + def signature_header(self, output: WriteableBuffer) -> int: ... + def sign_block(self, block: ReadableBuffer, output: WriteableBuffer) -> int: ... + def apply_delta_data(self, data: ReadableBuffer, read: Callable[[int, WriteableBuffer], int], write: Callable[[ReadableBuffer], None]) -> None: ... def finish_delta_data(self) -> None: ... @property @@ -40,9 +39,9 @@ class Patcher: class Differ: - def add_signature_data(self, data: ReadOnlyBuffer) -> None: ... + def add_signature_data(self, data: ReadableBuffer) -> None: ... def finish_signature_data(self) -> None: ... - def next_op(self, read: Callable[[WriteBuffer], int], write: Callable[[ReadOnlyBuffer], None]) -> bool: ... + def next_op(self, read: Callable[[WriteableBuffer], int], write: Callable[[ReadableBuffer], None]) -> bool: ... -def parse_ftc(x: Union[str, ReadOnlyBuffer], callback: Callable[[memoryview, memoryview], None]) -> None: ... +def parse_ftc(x: Union[str, ReadableBuffer], callback: Callable[[memoryview, memoryview], None]) -> None: ... diff --git a/kitty/data-types.c b/kitty/data-types.c index e3d98851e..7f4b3e2e8 100644 --- a/kitty/data-types.c +++ b/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}, }, diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi index e2ef0f170..ffc038b7f 100644 --- a/kitty/fast_data_types.pyi +++ b/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: ... diff --git a/kitty/file_transmission.py b/kitty/file_transmission.py index db4486751..ed49749ae 100644 --- a/kitty/file_transmission.py +++ b/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) diff --git a/kitty/typing.py b/kitty/typing.py index 1d35f044b..b667ceb75 100644 --- a/kitty/typing.py +++ 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 diff --git a/kitty/typing.pyi b/kitty/typing.pyi index 5f9c224b0..e183458d0 100644 --- a/kitty/typing.pyi +++ b/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', )