From 23bb2e1b67814e8c4697b3de0399c2d6f89379e4 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Fri, 3 Nov 2023 17:27:35 +0530 Subject: [PATCH] Fast function to replace c0 codes --- kitty/data-types.c | 71 +++++++++++++++++++++++++++++++++++++++ kitty/fast_data_types.pyi | 17 +++------- kitty_tests/datatypes.py | 12 +++++++ 3 files changed, 88 insertions(+), 12 deletions(-) diff --git a/kitty/data-types.c b/kitty/data-types.c index 20ab0a2b4..c86fb8d05 100644 --- a/kitty/data-types.c +++ b/kitty/data-types.c @@ -13,6 +13,7 @@ #endif #include "data-types.h" +#include "charsets.h" #include "base64.h" #include #include @@ -320,6 +321,75 @@ expand_ansi_c_escapes(PyObject *self UNUSED, PyObject *src) { return ans; } +START_ALLOW_CASE_RANGE +#define C0_EXCEPT_NL_AND_SPACE 0x0 ... 0x9: case 0xb ... 0x1f: case 0x7f +static PyObject* +c0_replace_bytes(const char *input_data, Py_ssize_t input_sz) { + RAII_PyObject(ans, PyBytes_FromStringAndSize(NULL, input_sz * 3)); + if (!ans) return NULL; + char *output = PyBytes_AS_STRING(ans); + char buf[4]; + Py_ssize_t j = 0; + for (Py_ssize_t i = 0; i < input_sz; i++) { + const char x = input_data[i]; + switch (x) { + case C0_EXCEPT_NL_AND_SPACE: { + const uint32_t ch = 0x2400 + x; + const unsigned sz = encode_utf8(ch, buf); + for (unsigned c = 0; c < sz; c++, j++) output[j] = buf[c]; + } break; + default: + output[j++] = x; break; + } + } + if (_PyBytes_Resize(&ans, j) != 0) return NULL; + Py_INCREF(ans); + return ans; +} + +static PyObject* +c0_replace_unicode(PyObject *input) { + RAII_PyObject(ans, PyUnicode_New(PyUnicode_GET_LENGTH(input), 1114111)); + if (!ans) return NULL; + void *input_data = PyUnicode_DATA(input); + int input_kind = PyUnicode_KIND(input); + void *output_data = PyUnicode_DATA(ans); + int output_kind = PyUnicode_KIND(ans); + Py_UCS4 maxchar = 0; + bool changed = false; + for (Py_ssize_t i = 0; i < PyUnicode_GET_LENGTH(input); i++) { + Py_UCS4 ch = PyUnicode_READ(input_kind, input_data, i); + switch(ch) { case C0_EXCEPT_NL_AND_SPACE: ch += 0x2400; changed = true; } + if (ch > maxchar) maxchar = ch; + PyUnicode_WRITE(output_kind, output_data, i, ch); + } + if (!changed) { Py_INCREF(input); return input; } + if (maxchar > 65535) { Py_INCREF(ans); return ans; } + RAII_PyObject(ans2, PyUnicode_New(PyUnicode_GET_LENGTH(ans), maxchar)); + if (!ans2) return NULL; + if (PyUnicode_CopyCharacters(ans2, 0, ans, 0, PyUnicode_GET_LENGTH(ans)) == -1) return NULL; + Py_INCREF(ans2); return ans2; +} +END_ALLOW_CASE_RANGE + +static PyObject* +replace_c0_codes_except_for_newline_and_space(PyObject *self UNUSED, PyObject *obj) { + if (PyUnicode_Check(obj)) { + return c0_replace_unicode(obj); + } else if (PyBytes_Check(obj)) { + return c0_replace_bytes(PyBytes_AS_STRING(obj), PyBytes_GET_SIZE(obj)); + } else if (PyMemoryView_Check(obj)) { + Py_buffer *buf = PyMemoryView_GET_BUFFER(obj); + return c0_replace_bytes(buf->buf, buf->len); + } else if (PyByteArray_Check(obj)) { + return c0_replace_bytes(PyByteArray_AS_STRING(obj), PyByteArray_GET_SIZE(obj)); + } else { + PyErr_SetString(PyExc_TypeError, "Input must be bytes, memoryview, bytearray or unicode"); + return NULL; + } +} + + static PyObject* find_in_memoryview(PyObject *self UNUSED, PyObject *args) { const char *buf; Py_ssize_t sz; @@ -332,6 +402,7 @@ find_in_memoryview(PyObject *self UNUSED, PyObject *args) { } static PyMethodDef module_methods[] = { + METHODB(replace_c0_codes_except_for_newline_and_space, METH_O), {"wcwidth", (PyCFunction)wcwidth_wrap, METH_O, ""}, {"expand_ansi_c_escapes", (PyCFunction)expand_ansi_c_escapes, METH_O, ""}, {"get_docs_ref_map", (PyCFunction)get_docs_ref_map, METH_NOARGS, ""}, diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi index da09b9131..54a055abf 100644 --- a/kitty/fast_data_types.pyi +++ b/kitty/fast_data_types.pyi @@ -1,17 +1,6 @@ import termios from ctypes import Array, c_ubyte -from typing import ( - Any, - Callable, - Dict, - Iterator, - List, - NewType, - Optional, - Tuple, - TypedDict, - Union, -) +from typing import Any, Callable, Dict, Iterator, List, NewType, Optional, Tuple, TypedDict, Union, overload from kitty.boss import Boss from kitty.fonts import FontFeature @@ -1567,3 +1556,7 @@ def cocoa_clear_global_shortcuts() -> None: ... def update_pointer_shape(os_window_id: int) -> None: ... def os_window_focus_counters() -> Dict[int, int]: ... def find_in_memoryview(buf: Union[bytes, memoryview, bytearray], chr: int) -> int: ... +@overload +def replace_c0_codes_except_for_newline_and_space(text: str) -> str:... +@overload +def replace_c0_codes_except_for_newline_and_space(text: Union[bytes, memoryview, bytearray]) -> bytes:... diff --git a/kitty_tests/datatypes.py b/kitty_tests/datatypes.py index 50db63a87..cb17b01e3 100644 --- a/kitty_tests/datatypes.py +++ b/kitty_tests/datatypes.py @@ -13,6 +13,7 @@ from kitty.fast_data_types import ( LineBuf, expand_ansi_c_escapes, parse_input_from_terminal, + replace_c0_codes_except_for_newline_and_space, strip_csi, truncate_point_for_length, wcswidth, @@ -37,6 +38,17 @@ def create_lbuf(*lines): class TestDataTypes(BaseTest): + + def test_replace_c0_codes(self): + def t(x: str, expected: str): + q = replace_c0_codes_except_for_newline_and_space(x) + self.ae(expected, q) + q = replace_c0_codes_except_for_newline_and_space(x.encode('utf-8')) + self.ae(expected.encode('utf-8'), q) + t('abc', 'abc') + t('a\0\x01b\x03\x04\t\rc', 'a\u2400\u2401b\u2403\u2404\u2409\u240dc') + t('a\0\x01😸\x03\x04\t\rc', 'a\u2400\u2401😸\u2403\u2404\u2409\u240dc') + def test_to_color(self): for x in 'xxx #12 #1234 rgb:a/b'.split(): self.assertIsNone(to_color(x))