From 182bd9c7b329973fc45d6d133ec0748488d66038 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Fri, 17 Jan 2025 19:28:55 +0530
Subject: [PATCH] Pass the basic colors of the underlying window as an env var
 when running a UI kitten

Allows UI kittens to choose colors without needing a roundtrip to query
the terminal for colors
---
 kitty/boss.py             |  2 ++
 kitty/colors.c            | 52 ++++++++++++++++++++++++++++-----------
 kitty/fast_data_types.pyi |  4 +--
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/kitty/boss.py b/kitty/boss.py
index 96ea2023e51..e6ddfbb39cf 100644
--- a/kitty/boss.py
+++ b/kitty/boss.py
@@ -1976,6 +1976,8 @@ def run_kitten_with_metadata(
                 cmd = [kitten_exe(), kitten]
                 env['KITTEN_RUNNING_AS_UI'] = '1'
                 env['KITTY_CONFIG_DIRECTORY'] = config_dir
+                if w is not None:
+                    env['KITTY_BASIC_COLORS'] = json.dumps(w.screen.color_profile.basic_colors())
             else:
                 cmd = [kitty_exe(), '+runpy', 'from kittens.runner import main; main()']
                 env['PYTHONWARNINGS'] = 'ignore'
diff --git a/kitty/colors.c b/kitty/colors.c
index 6fbcabafd45..aa105ff98ef 100644
--- a/kitty/colors.c
+++ b/kitty/colors.c
@@ -305,32 +305,55 @@ colorprofile_to_color_with_fallback(ColorProfile *self, DynamicColor entry, Dyna
 }
 static Color* alloc_color(unsigned char r, unsigned char g, unsigned char b, unsigned a);
 
-static PyObject*
-as_dict(ColorProfile *self, PyObject *args UNUSED) {
-#define as_dict_doc "Return all colors as a dictionary of color_name to integer or None (names are the same as used in kitty.conf)"
-    RAII_PyObject(ans, PyDict_New());
-    if (ans == NULL) return PyErr_NoMemory();
-    for (unsigned i = 0; i < arraysz(self->color_table); i++) {
-        static char buf[32] = {0};
-        snprintf(buf, sizeof(buf) - 1, "color%u", i);
+static bool
+colortable_colors_into_dict(ColorProfile *self, unsigned start, unsigned limit, PyObject *ans) {
+    static char buf[32] = {'c', 'o', 'l', 'o', 'r', 0};
+    for (unsigned i = start; i < limit; i++) {
+        snprintf(buf + 5, sizeof(buf) - 6, "%u", i);
         PyObject *val = PyLong_FromUnsignedLong(self->color_table[i]);
-        if (!val) { return PyErr_NoMemory(); }
+        if (!val) return false;
         int ret = PyDict_SetItemString(ans, buf, val);
-        Py_CLEAR(val);
-        if (ret != 0) { return NULL; }
+        Py_DECREF(val);
+        if (ret != 0) return false;
     }
+    return true;
+}
+
+static PyObject*
+basic_colors(ColorProfile *self, PyObject *args UNUSED) {
+#define basic_colors_doc "Return the basic colors as a dictionary of color_name to integer or None (names are the same as used in kitty.conf)"
+    RAII_PyObject(ans, PyDict_New()); if (ans == NULL) return NULL;
+    if (!colortable_colors_into_dict(self, 0, 16, ans)) return NULL;
+
+#define D(attr, name) { \
+    unsigned long c = colorprofile_to_color(self, self->overridden.attr, self->configured.attr).rgb; \
+    PyObject *val = PyLong_FromUnsignedLong(c); if (!val) return NULL; \
+    int ret = PyDict_SetItemString(ans, #name, val); Py_DECREF(val); \
+    if (ret != 0) return NULL; \
+}
+
+    D(default_fg, foreground); D(default_bg, background);
+#undef D
+    return Py_NewRef(ans);
+}
+
+static PyObject*
+as_dict(ColorProfile *self, PyObject *args UNUSED) {
+#define as_dict_doc "Return all colors as a dictionary of color_name to integer or None (names are the same as used in kitty.conf)"
+    RAII_PyObject(ans, PyDict_New()); if (ans == NULL) return NULL;
+    if (!colortable_colors_into_dict(self, 0, arraysz(self->color_table), ans)) return NULL;
 #define D(attr, name) { \
     if (self->overridden.attr.type != COLOR_NOT_SET) { \
         int ret; PyObject *val; \
         if (self->overridden.attr.type == COLOR_IS_SPECIAL) { \
-            val = Py_None; Py_INCREF(val); \
+            val = Py_NewRef(Py_None); \
         } else { \
-            color_type c = colorprofile_to_color(self, self->overridden.attr, self->configured.attr).rgb; \
+            unsigned long c = colorprofile_to_color(self, self->overridden.attr, self->configured.attr).rgb; \
             val = PyLong_FromUnsignedLong(c); \
         } \
         if (!val) { return NULL; } \
         ret = PyDict_SetItemString(ans, #name, val); \
-        Py_CLEAR(val); \
+        Py_DECREF(val); \
         if (ret != 0) { return NULL; } \
     }}
     D(default_fg, foreground); D(default_bg, background);
@@ -597,6 +620,7 @@ set_transparent_background_color(ColorProfile *self, PyObject *const *args, Py_s
 static PyMethodDef cp_methods[] = {
     METHOD(reset_color_table, METH_NOARGS)
     METHOD(as_dict, METH_NOARGS)
+    METHOD(basic_colors, METH_NOARGS)
     METHOD(color_table_address, METH_NOARGS)
     METHOD(as_color, METH_O)
     METHOD(reset_color, METH_O)
diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi
index 94c4743675f..c91aa35c25d 100644
--- a/kitty/fast_data_types.pyi
+++ b/kitty/fast_data_types.pyi
@@ -821,8 +821,8 @@ class ColorProfile:
 
     def __init__(self, opts: Optional[Options] = None): ...
 
-    def as_dict(self) -> Dict[str, int | None | tuple[tuple[Color, float], ...]]:
-        pass
+    def as_dict(self) -> Dict[str, int | None | tuple[tuple[Color, float], ...]]: ...
+    def basic_colors(self) -> Dict[str, int | None | tuple[tuple[Color, float], ...]]: ...
 
     def as_color(self, val: int) -> Optional[Color]:
         pass