Fast function to replace c0 codes

kitty/data-types.c

@@ -13,6 +13,7 @@
 #endif
 
 #include "data-types.h"
+#include "charsets.h"
 #include "base64.h"
 #include <sys/socket.h>
 #include <sys/types.h>
@@ -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, ""},

kitty/fast_data_types.pyi

@@ -1,17 +1,6 @@
 import termios
 from ctypes import Array, c_ubyte
-from typing import (
+from typing import Any, Callable, Dict, Iterator, List, NewType, Optional, Tuple, TypedDict, Union, overload
-    Any,
-    Callable,
-    Dict,
-    Iterator,
-    List,
-    NewType,
-    Optional,
-    Tuple,
-    TypedDict,
-    Union,
-)
 
 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:...

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))