Allow sending escape code to child in chunks

kitty/child-monitor.c

@@ -228,56 +228,92 @@ add_child(ChildMonitor *self, PyObject *args) {
     Py_RETURN_NONE;
 }
 
+#define schedule_write_to_child_generic(id, num, va_start, get_next_arg, va_end) \
+    ChildMonitor *self = the_monitor; \
+    bool found = false; \
+    const char *data; \
+    size_t szval, sz = 0; \
+    va_start(ap, num); \
+    for (unsigned int i = 0; i < num; i++) { \
+        get_next_arg(ap); \
+        sz += szval; \
+    } \
+    va_end(ap); \
+    children_mutex(lock); \
+    for (size_t i = 0; i < self->count; i++) { \
+        if (children[i].id == id) { \
+            Screen *screen = children[i].screen; \
+            screen_mutex(lock, write); \
+            size_t space_left = screen->write_buf_sz - screen->write_buf_used; \
+            if (space_left < sz) { \
+                if (screen->write_buf_used + sz > 100 * 1024 * 1024) { \
+                    log_error("Too much data being sent to child with id: %lu, ignoring it", id); \
+                    screen_mutex(unlock, write); \
+                    break; \
+                } \
+                screen->write_buf_sz = screen->write_buf_used + sz; \
+                screen->write_buf = PyMem_RawRealloc(screen->write_buf, screen->write_buf_sz); \
+                if (screen->write_buf == NULL) { fatal("Out of memory."); } \
+            } \
+            found = true; \
+            va_start(ap, num); \
+            for (unsigned int i = 0; i < num; i++) { \
+                get_next_arg(ap); \
+                memcpy(screen->write_buf + screen->write_buf_used, data, szval); \
+                screen->write_buf_used += szval; \
+            } \
+            va_end(ap); \
+            if (screen->write_buf_sz > BUFSIZ && screen->write_buf_used < BUFSIZ) { \
+                screen->write_buf_sz = BUFSIZ; \
+                screen->write_buf = PyMem_RawRealloc(screen->write_buf, screen->write_buf_sz); \
+                if (screen->write_buf == NULL) { fatal("Out of memory."); } \
+            } \
+            if (screen->write_buf_used) wakeup_io_loop(self, false); \
+            screen_mutex(unlock, write); \
+            break; \
+        } \
+    } \
+    children_mutex(unlock); \
+    return found;
+
 bool
 schedule_write_to_child(unsigned long id, unsigned int num, ...) {
-    ChildMonitor *self = the_monitor;
-    bool found = false;
-    const char *data;
-    size_t sz = 0;
     va_list ap;
-    va_start(ap, num);
+#define get_next_arg(ap) data = va_arg(ap, const char*); szval = va_arg(ap, size_t);
-    for (unsigned int i = 0; i < num; i++) {
+    schedule_write_to_child_generic(id, num, va_start, get_next_arg, va_end);
-        va_arg(ap, const char*);
+#undef get_next_arg
-        sz += va_arg(ap, size_t);
+}
-    }
+
-    va_end(ap);
+bool
-    children_mutex(lock);
+schedule_write_to_child_python(unsigned long id, const char *prefix, PyObject *ap, const char *suffix) {
-    for (size_t i = 0; i < self->count; i++) {
+    if (!PyTuple_Check(ap)) return false;
-        if (children[i].id == id) {
+    bool has_prefix = prefix && prefix[0], has_suffix = suffix && suffix[0];
-            Screen *screen = children[i].screen;
+    const size_t extra = (has_prefix ? 1 : 0) + (has_suffix ? 1 : 0);
-            screen_mutex(lock, write);
+    size_t num = PyTuple_GET_SIZE(ap) + extra;
-            size_t space_left = screen->write_buf_sz - screen->write_buf_used;
+    Py_ssize_t pidx;
-            if (space_left < sz) {
+#define py_start(ap, num) pidx = 0;
-                if (screen->write_buf_used + sz > 100 * 1024 * 1024) {
+#define py_end(ap) pidx = 0;
-                    log_error("Too much data being sent to child with id: %lu, ignoring it", id);
+#define get_next_arg(ap) { \
-                    screen_mutex(unlock, write);
+    if (pidx == 0 && has_prefix) { data = prefix; szval = strlen(prefix); } \
-                    break;
+    else { \
-                }
+        size_t pidxf = pidx++; \
-                screen->write_buf_sz = screen->write_buf_used + sz;
+        if (has_prefix) pidxf--; \
-                screen->write_buf = PyMem_RawRealloc(screen->write_buf, screen->write_buf_sz);
+        if (has_suffix && pidxf >= (size_t)PyBytes_GET_SIZE(ap)) { data = suffix; szval = strlen(suffix); } \
-                if (screen->write_buf == NULL) { fatal("Out of memory."); }
+        else { \
-            }
+            PyObject *t = PyTuple_GET_ITEM(ap, pidxf); \
-            found = true;
+            if (PyBytes_Check(t)) { data = PyBytes_AS_STRING(t); szval = PyBytes_GET_SIZE(t); } \
-            va_start(ap, num);
+            else { \
-            for (unsigned int i = 0; i < num; i++) {
+                Py_ssize_t usz; \
-                data = va_arg(ap, const char*);
+                data = PyUnicode_AsUTF8AndSize(t, &usz); szval = usz; \
-                size_t dsz = va_arg(ap, size_t);
+                if (!data) fatal("Failed to convert object to bytes in schedule_write_to_child_python"); \
-                memcpy(screen->write_buf + screen->write_buf_used, data, dsz);
+            } \
-                screen->write_buf_used += dsz;
+        } \
-            }
+    } \
-            va_end(ap);
+}
-            if (screen->write_buf_sz > BUFSIZ && screen->write_buf_used < BUFSIZ) {
+    schedule_write_to_child_generic(id, num, py_start, get_next_arg, py_end);
-                screen->write_buf_sz = BUFSIZ;
+#undef py_start
-                screen->write_buf = PyMem_RawRealloc(screen->write_buf, screen->write_buf_sz);
+#undef py_end
-                if (screen->write_buf == NULL) { fatal("Out of memory."); }
+#undef get_next_arg
-            }
-            if (screen->write_buf_used) wakeup_io_loop(self, false);
-            screen_mutex(unlock, write);
-            break;
-        }
-    }
-    children_mutex(unlock);
-    return found;
 }
 
 static PyObject *

kitty/data-types.h

@@ -326,6 +326,7 @@ const char* cursor_as_sgr(const Cursor *);
 
 PyObject* cm_thread_write(PyObject *self, PyObject *args);
 bool schedule_write_to_child(unsigned long id, unsigned int num, ...);
+bool schedule_write_to_child_python(unsigned long id, const char *prefix, PyObject* tuple_of_str_or_bytes, const char *suffix);
 bool set_iutf8(int, bool);
 
 color_type colorprofile_to_color(ColorProfile *self, color_type entry, color_type defval);

kitty/fast_data_types.pyi

@@ -993,7 +993,7 @@ class Screen:
     def resize(self, width: int, height: int) -> None:
         pass
 
-    def send_escape_code_to_child(self, code: int, text: Union[str, bytes]) -> bool:
+    def send_escape_code_to_child(self, code: int, text: Union[str, bytes, Tuple[Union[str, bytes], ...]]) -> bool:
         pass
 
     def reset_callbacks(self) -> None:

kitty/screen.c

@@ -780,29 +780,35 @@ write_to_child(Screen *self, const char *data, size_t sz) {
     return written;
 }
 
-bool
+static void
-write_escape_code_to_child(Screen *self, unsigned char which, const char *data) {
+get_prefix_and_suffix_for_escape_code(const Screen *self, unsigned char which, const char ** prefix, const char ** suffix) {
-    bool written = false;
+    *suffix = self->modes.eight_bit_controls ? "\x9c" : "\033\\";
-    const char *prefix, *suffix = self->modes.eight_bit_controls ? "\x9c" : "\033\\";
     switch(which) {
         case DCS:
-            prefix = self->modes.eight_bit_controls ? "\x90" : "\033P";
+            *prefix = self->modes.eight_bit_controls ? "\x90" : "\033P";
             break;
         case CSI:
-            prefix = self->modes.eight_bit_controls ? "\x9b" : "\033["; suffix = "";
+            *prefix = self->modes.eight_bit_controls ? "\x9b" : "\033["; *suffix = "";
             break;
         case OSC:
-            prefix = self->modes.eight_bit_controls ? "\x9d" : "\033]";
+            *prefix = self->modes.eight_bit_controls ? "\x9d" : "\033]";
             break;
         case PM:
-            prefix = self->modes.eight_bit_controls ? "\x9e" : "\033^";
+            *prefix = self->modes.eight_bit_controls ? "\x9e" : "\033^";
             break;
         case APC:
-            prefix = self->modes.eight_bit_controls ? "\x9f" : "\033_";
+            *prefix = self->modes.eight_bit_controls ? "\x9f" : "\033_";
             break;
         default:
             fatal("Unknown escape code to write: %u", which);
     }
+}
+
+bool
+write_escape_code_to_child(Screen *self, unsigned char which, const char *data) {
+    bool written = false;
+    const char *prefix, *suffix;
+    get_prefix_and_suffix_for_escape_code(self, which, &prefix, &suffix);
     if (self->window_id) {
         if (suffix[0]) {
             written = schedule_write_to_child(self->window_id, 3, prefix, strlen(prefix), data, strlen(data), suffix, strlen(suffix));
@@ -818,6 +824,28 @@ write_escape_code_to_child(Screen *self, unsigned char which, const char *data)
     return written;
 }
 
+static bool
+write_escape_code_to_child_python(Screen *self, unsigned char which, PyObject *data) {
+    bool written = false;
+    const char *prefix, *suffix;
+    get_prefix_and_suffix_for_escape_code(self, which, &prefix, &suffix);
+    if (self->window_id) written = schedule_write_to_child_python(self->window_id, prefix, data, suffix);
+    if (self->test_child != Py_None) {
+        write_to_test_child(self, prefix, strlen(prefix));
+        for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(data); i++) {
+            PyObject *t = PyTuple_GET_ITEM(data, i);
+            if (PyBytes_Check(t)) write_to_test_child(self, PyBytes_AS_STRING(t), PyBytes_GET_SIZE(t));
+            else {
+                Py_ssize_t sz;
+                const char *d = PyUnicode_AsUTF8AndSize(t, &sz);
+                if (d) write_to_test_child(self, d, sz);
+            }
+        }
+        if (suffix[0]) write_to_test_child(self, suffix, strlen(suffix));
+    }
+    return written;
+}
+
 static bool
 cursor_within_margins(Screen *self) {
     return self->margin_top <= self->cursor->y && self->cursor->y <= self->margin_bottom;
@@ -3234,11 +3262,18 @@ toggle_alt_screen(Screen *self, PyObject *a UNUSED) {
 static PyObject*
 send_escape_code_to_child(Screen *self, PyObject *args) {
     int code;
-    char *text;
+    PyObject *O;
-    Py_ssize_t sz;
+    if (!PyArg_ParseTuple(args, "iO", &code, &O)) return NULL;
-    if (!PyArg_ParseTuple(args, "is#", &code, &text, &sz)) return NULL;
+    bool written = false;
-    if (write_escape_code_to_child(self, code, text)) Py_RETURN_TRUE;
+    if (PyBytes_Check(O)) written = write_escape_code_to_child(self, code, PyBytes_AS_STRING(O));
-    Py_RETURN_FALSE;
+    else if (PyUnicode_Check(O)) {
+        const char *t = PyUnicode_AsUTF8(O);
+        if (t) written = write_escape_code_to_child(self, code, t);
+        else return NULL;
+    } else if (PyTuple_Check(O)) written = write_escape_code_to_child_python(self, code, O);
+    else PyErr_SetString(PyExc_TypeError, "escape code must be str, bytes or tuple");
+    if (PyErr_Occurred()) return NULL;
+    if (written) { Py_RETURN_TRUE; } else { Py_RETURN_FALSE; }
 }
 
 static void