From 5d120a2f36d8393928d40a0d4344316142e91028 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 15 Feb 2022 08:57:17 +0530
Subject: [PATCH] Output prompt marking when serializing to ANSI

---
 kitty/line.c          | 25 +++++++++++++++++++++++++
 kitty_tests/screen.py |  5 +++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/kitty/line.c b/kitty/line.c
index db2942093..daa4726cf 100644
--- a/kitty/line.c
+++ b/kitty/line.c
@@ -301,6 +301,15 @@ write_hyperlink(hyperlink_id_type hid, ANSIBuf *output) {
 #undef W
 }
 
+static void
+write_mark(const char *mark, ANSIBuf *output) {
+#define W(c) output->buf[output->len++] = c
+    W(0x1b); W(']'); W('1'); W('3'); W('3'); W(';');
+    for (size_t i = 0; mark[i] != 0 && i < 32; i++) W(mark[i]);
+    W(0x1b); W('\\');
+#undef W
+
+}
 
 bool
 line_as_ansi(Line *self, ANSIBuf *output, const GPUCell** prev_cell, index_type start_at, index_type stop_before, char_type prefix_char) {
@@ -308,11 +317,26 @@ line_as_ansi(Line *self, ANSIBuf *output, const GPUCell** prev_cell, index_type
 #define WRITE_SGR(val) { ENSURE_SPACE(128); escape_code_written = true; write_sgr(val, output); }
 #define WRITE_CH(val) { ENSURE_SPACE(1); output->buf[output->len++] = val; }
 #define WRITE_HYPERLINK(val) { ENSURE_SPACE(2256); escape_code_written = true; write_hyperlink(val, output); }
+#define WRITE_MARK(val) { ENSURE_SPACE(64); escape_code_written = true; write_mark(val, output); }
     bool escape_code_written = false;
     output->len = 0;
     index_type limit = MIN(stop_before, xlimit_for_line(self));
     char_type previous_width = 0;
     if (prefix_char) { WRITE_CH(prefix_char); previous_width = wcwidth_std(prefix_char); }
+
+    switch (self->attrs.prompt_kind) {
+        case UNKNOWN_PROMPT_KIND:
+            break;
+        case PROMPT_START:
+            WRITE_MARK("A");
+            break;
+        case SECONDARY_PROMPT:
+            WRITE_MARK("A;k=s");
+            break;
+        case OUTPUT_START:
+            WRITE_MARK("C");
+            break;
+    }
     if (limit <= start_at) return escape_code_written;
 
     static const GPUCell blank_cell = { 0 };
@@ -362,6 +386,7 @@ line_as_ansi(Line *self, ANSIBuf *output, const GPUCell** prev_cell, index_type
 #undef WRITE_CH
 #undef ENSURE_SPACE
 #undef WRITE_HYPERLINK
+#undef WRITE_MARK
 }
 
 static PyObject*
diff --git a/kitty_tests/screen.py b/kitty_tests/screen.py
index 7fb7c7ffa..e42cd35ea 100644
--- a/kitty_tests/screen.py
+++ b/kitty_tests/screen.py
@@ -996,9 +996,9 @@ def mark_output():
         self.assertTrue(s.scroll_to_prompt())
         self.ae(str(s.visual_line(0)), '$ 1')
 
-        def lco():
+        def lco(as_ansi=False):
             a = []
-            s.cmd_output(0, a.append)
+            s.cmd_output(0, a.append, as_ansi)
             return ''.join(a)
 
         def fco():
@@ -1025,6 +1025,7 @@ def lvco():
         mark_prompt(), s.draw('$ 1')
         self.ae(fco(), 'abcd\n12')
         self.ae(lco(), 'abcd\n12')
+        self.ae(lco(as_ansi=True), '\x1b[m\x1b]133;C\x1b\\abcd\n\x1b[m12')
 
         def draw_prompt(x):
             mark_prompt(), s.draw(f'$ {x}'), s.carriage_return(), s.index()