From c81ac668da1c6c4b84695296dee671bdeb871cdd Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 31 Oct 2023 18:14:49 +0530
Subject: [PATCH] Use a single code path for tests and live VT parsing

---
 kitty/screen.c          | 20 ++++++++++++++++++++
 kitty/vt-parser.c       | 27 ---------------------------
 kitty_tests/__init__.py |  4 ++--
 kitty_tests/parser.py   |  6 ++++--
 4 files changed, 26 insertions(+), 31 deletions(-)

diff --git a/kitty/screen.c b/kitty/screen.c
index 27b45135c..3967318f5 100644
--- a/kitty/screen.c
+++ b/kitty/screen.c
@@ -4454,7 +4454,27 @@ WRAP0(bell)
 #define MND(name, args) {#name, (PyCFunction)name, args, #name},
 #define MODEFUNC(name) MND(name, METH_NOARGS) MND(set_##name, METH_O)
 
+static PyObject*
+test_write_data(Screen *screen, PyObject *args) {
+    const uint8_t *data; Py_ssize_t sz;
+    PyObject *dump_callback = NULL;
+    if (!PyArg_ParseTuple(args, "y#|O", &data, &sz, &dump_callback)) return NULL;
+
+    monotonic_t now = monotonic();
+    while (sz) {
+        size_t s = MIN(sz, READ_BUF_SZ);
+        memcpy(screen->read_buf, data, s);
+        screen->read_buf_sz = s;
+        data += s; sz -= s;
+        if (dump_callback) parse_worker_dump(screen, dump_callback, now);
+        else parse_worker(screen, dump_callback, now);
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyMethodDef methods[] = {
+    METHODB(test_write_data, METH_VARARGS),
     MND(line_edge_colors, METH_NOARGS)
     MND(line, METH_O)
     MND(dump_lines_with_attrs, METH_O)
diff --git a/kitty/vt-parser.c b/kitty/vt-parser.c
index 8cf51b544..768d2db2e 100644
--- a/kitty/vt-parser.c
+++ b/kitty/vt-parser.c
@@ -1625,20 +1625,9 @@ run_worker(Screen *screen, PyObject *dump_callback, monotonic_t now) {
 }
 
 #ifdef DUMP_COMMANDS
-void
-parse_vt_dump(Parser *p) {
-    do_parse_vt((PS*)p->state);
-}
-
 void
 parse_worker_dump(Screen *screen, PyObject *dump_callback, monotonic_t now) { run_worker(screen, dump_callback, now); }
 #else
-static void
-parse_vt(Parser *p) {
-    do_parse_vt((PS*)p->state);
-}
-extern void parse_vt_dump(Parser *p);
-
 void
 parse_worker(Screen *screen, PyObject *dump_callback, monotonic_t now) { run_worker(screen, dump_callback, now); }
 #endif
@@ -1681,23 +1670,7 @@ reset_vt_parser(Parser *self) {
 
 extern PyTypeObject Screen_Type;
 
-static PyObject*
-py_parse(Parser *p, PyObject *args) {
-    PS *self = (PS*)p->state;
-    const uint8_t *data; Py_ssize_t sz;
-    PyObject *dump_callback = NULL;
-    if (!PyArg_ParseTuple(args, "O!y#|O", &Screen_Type, &self->screen, &data, &sz, &dump_callback)) return NULL;
-
-    self->input_data = data; self->input_sz = sz; self->dump_callback = dump_callback;
-    self->input_pos = 0; self->now = monotonic();
-    if (dump_callback) parse_vt_dump(p); else parse_vt(p);
-    self->input_data = NULL; self->input_sz = 0; self->dump_callback = NULL; self->screen = NULL;
-
-    Py_RETURN_NONE;
-}
-
 static PyMethodDef methods[] = {
-    {"parse_bytes", (PyCFunction)py_parse, METH_VARARGS, ""},
     {NULL},
 };
 
diff --git a/kitty_tests/__init__.py b/kitty_tests/__init__.py
index 023c72674..ffadad336 100644
--- a/kitty_tests/__init__.py
+++ b/kitty_tests/__init__.py
@@ -25,7 +25,7 @@ from kitty.window import process_remote_print, process_title_from_child
 
 
 def parse_bytes(screen, data):
-    return screen.vt_parser.parse_bytes(screen, data)
+    return screen.test_write_data(data)
 
 
 class Callbacks:
@@ -324,7 +324,7 @@ class PTY:
                 break
             bytes_read += len(data)
             self.received_bytes += data
-            self.screen.vt_parser.parse_bytes(self.screen, data)
+            self.screen.test_write_data(data)
         return bytes_read
 
     def wait_till(self, q, timeout=10):
diff --git a/kitty_tests/parser.py b/kitty_tests/parser.py
index 66aabbb8e..199f69673 100644
--- a/kitty_tests/parser.py
+++ b/kitty_tests/parser.py
@@ -23,6 +23,8 @@ class CmdDump(list):
     def __call__(self, *a):
         if a and isinstance(a[0], int):
             a = a[1:]
+        if a and a[0] == 'bytes':
+            return
         self.append(tuple(map(cnv, a)))
 
 
@@ -33,7 +35,7 @@ class TestParser(BaseTest):
         if isinstance(x, str):
             x = x.encode('utf-8')
         cmds = tuple(('draw', x) if isinstance(x, str) else tuple(map(cnv, x)) for x in cmds)
-        s.vt_parser.parse_bytes(s, x, cd)
+        s.test_write_data(x, cd)
         current = ''
         q = []
         for args in cd:
@@ -73,7 +75,7 @@ class TestParser(BaseTest):
         self.ae(str(s.line(1)), '6')
         self.ae(str(s.line(2)), ' 123')
         self.ae(str(s.line(3)), '45')
-        s.vt_parser.parse_bytes(s, b'\rabcde')
+        s.test_write_data(b'\rabcde')
         self.ae(str(s.line(3)), 'abcde')
         pb('\rßxyz1', ('screen_carriage_return',), 'ßxyz1')
         self.ae(str(s.line(3)), 'ßxyz1')