fix a bug

include/pocketpy/interpreter/modules.h

@@ -26,6 +26,7 @@ void pk__add_module_colorcvt();
 void pk__add_module_conio();
 void pk__add_module_lz4();
 void pk__add_module_pkpy();
+void pk__add_module_picoterm();
 
 #ifdef PK_BUILD_MODULE_CUTE_PNG
 void pk__add_module_cute_png();

include/pocketpy/vmath.h

@@ -42,4 +42,5 @@ typedef union c11_color32 {
         unsigned char a;
     };
     unsigned char data[4];
+    uint32_t u32;
 } c11_color32;

include/typings/array2d.pyi

@@ -1,5 +1,5 @@
 from typing import Callable, Literal, overload, Iterator, Self
-from vmath import vec2i
+from vmath import vec2i, color32
 
 Neighborhood = Literal['Moore', 'von Neumann']
 
@@ -35,6 +35,7 @@ class array2d_like[T]:
         """
 
     def render(self) -> str: ...
+    def render_with_color(self, fg: array2d_like[color32 | None], bg: array2d_like[color32 | None]) -> str: ...
 
     def all(self: array2d_like[bool]) -> bool: ...
     def any(self: array2d_like[bool]) -> bool: ...

include/typings/picoterm.pyi

@@ -0,0 +1,5 @@
+def enable_full_buffering_mode() -> None:
+    """Enable full buffering mode for ASCII drawings (32KB)."""
+
+def split_ansi_escaped_string(s: str) -> list[str]:
+    """Perform split on ANSI escaped string."""

include/typings/pkpy.pyi

@@ -20,9 +20,6 @@ def memory_usage() -> str:
 def is_user_defined_type(t: type) -> bool:
     """Check if a type is user-defined. This means the type was created by executing python `class` statement."""
 
-def enable_full_buffering_mode() -> None:
-    """Enable full buffering mode for ASCII drawings."""
-
 def currentvm() -> int:
     """Return the current VM index."""
 

src/interpreter/vm.c

@@ -265,6 +265,7 @@ void VM__ctor(VM* self) {
     pk__add_module_cute_png();  // optional
     pk__add_module_msgpack();   // optional
     pk__add_module_pkpy();
+    pk__add_module_picoterm();
 
     // add python builtins
     do {

src/modules/array2d.c

@@ -22,7 +22,7 @@ c11_array2d* c11_newarray2d(py_OutRef out, int n_cols, int n_rows) {
     ud->header.n_cols = n_cols;
     ud->header.n_rows = n_rows;
     ud->header.numel = numel;
-    ud->header.f_get = (py_Ref(*)(c11_array2d_like*, int, int))c11_array2d__get;
+    ud->header.f_get = (py_Ref (*)(c11_array2d_like*, int, int))c11_array2d__get;
     ud->header.f_set = (bool (*)(c11_array2d_like*, int, int, py_Ref))c11_array2d__set;
     ud->data = py_getslot(out, 0);
     return ud;
@@ -131,7 +131,10 @@ static bool array2d_like_render(int argc, py_Ref argv) {
     for(int j = 0; j < self->n_rows; j++) {
         for(int i = 0; i < self->n_cols; i++) {
             py_Ref item = self->f_get(self, i, j);
-            if(!py_str(item)) return false;
+            if(!py_str(item)) {
+                c11_sbuf__dtor(&buf);
+                return false;
+            }
             c11_sbuf__write_sv(&buf, py_tosv(py_retval()));
         }
         if(j < self->n_rows - 1) c11_sbuf__write_char(&buf, '\n');
@@ -140,6 +143,84 @@ static bool array2d_like_render(int argc, py_Ref argv) {
     return true;
 }
 
+void c11_color32_premult(c11_color32* color);
+
+static bool array2d_like_render_with_color(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(3);
+    c11_sbuf buf;
+    c11_sbuf__ctor(&buf);
+    c11_array2d_like* self = py_touserdata(argv);
+
+    if(!py_checkinstance(py_arg(1), tp_array2d_like)) return false;
+    if(!py_checkinstance(py_arg(2), tp_array2d_like)) return false;
+    c11_array2d_like* fg_colors = py_touserdata(py_arg(1));
+    c11_array2d_like* bg_colors = py_touserdata(py_arg(2));
+
+    c11_color32 curr_fg, curr_bg;
+    curr_fg.u32 = 0;
+    curr_bg.u32 = 0;
+
+    for(int j = 0; j < self->n_rows; j++) {
+        for(int i = 0; i < self->n_cols; i++) {
+            py_Ref item = self->f_get(self, i, j);
+            if(!py_str(item)) {
+                c11_sbuf__dtor(&buf);
+                return false;
+            }
+
+            py_Ref fg_item = fg_colors->f_get(fg_colors, i, j);
+            py_Ref bg_item = bg_colors->f_get(bg_colors, i, j);
+
+            c11_color32 new_fg, new_bg;
+            if(py_isnone(fg_item)) {
+                new_fg.u32 = 0;
+            } else {
+                if(!py_checktype(fg_item, tp_color32)) {
+                    c11_sbuf__dtor(&buf);
+                    return false;
+                }
+                new_fg = py_tocolor32(fg_item);
+            }
+            if(py_isnone(bg_item)) {
+                new_bg.u32 = 0;
+            } else {
+                if(!py_checktype(bg_item, tp_color32)) {
+                    c11_sbuf__dtor(&buf);
+                    return false;
+                }
+                new_bg = py_tocolor32(bg_item);
+            }
+
+            if(curr_fg.u32 != new_fg.u32 || curr_bg.u32 != new_bg.u32) {
+                if(curr_fg.u32 != 0 || curr_bg.u32 != 0) c11_sbuf__write_cstr(&buf, "\x1b[0m");
+                curr_fg = new_fg;
+                curr_bg = new_bg;
+                if(curr_fg.u32 != 0) {
+                    c11_color32_premult(&curr_fg);
+                    pk_sprintf(&buf, "\x1b[38;2;%d;%d;%dm", curr_fg.r, curr_fg.g, curr_fg.b);
+                }
+                if(curr_bg.u32 != 0) {
+                    c11_color32_premult(&curr_bg);
+                    pk_sprintf(&buf, "\x1b[48;2;%d;%d;%dm", curr_bg.r, curr_bg.g, curr_bg.b);
+                }
+            }
+
+            c11_sbuf__write_sv(&buf, py_tosv(py_retval()));
+        }
+        if(j < self->n_rows - 1) {
+            if(curr_fg.u32 != 0 || curr_bg.u32 != 0) {
+                curr_fg.u32 = 0;
+                curr_bg.u32 = 0;
+                c11_sbuf__write_cstr(&buf, "\x1b[0m\n");
+            } else {
+                c11_sbuf__write_char(&buf, '\n');
+            }
+        }
+    }
+    c11_sbuf__py_submit(&buf, py_retval());
+    return true;
+}
+
 static bool array2d_like_all(int argc, py_Ref argv) {
     PY_CHECK_ARGC(1);
     c11_array2d_like* self = py_touserdata(argv);
@@ -409,7 +490,7 @@ static c11_array2d_view* _array2d_view__new(py_OutRef out,
     res->header.n_cols = width;
     res->header.n_rows = height;
     res->header.numel = width * height;
-    res->header.f_get = (py_Ref(*)(c11_array2d_like*, int, int))c11_array2d_view__get;
+    res->header.f_get = (py_Ref (*)(c11_array2d_like*, int, int))c11_array2d_view__get;
     res->header.f_set = (bool (*)(c11_array2d_like*, int, int, py_Ref))c11_array2d_view__set;
     res->origin.x = start_col;
     res->origin.y = start_row;
@@ -427,7 +508,7 @@ static bool _array2d_view(py_OutRef out,
     c11_array2d_view* res = _array2d_view__new(out, keepalive, start_col, start_row, width, height);
     if(res == NULL) return false;
     res->ctx = array;
-    res->f_get = (py_Ref(*)(void*, int, int))array->f_get;
+    res->f_get = (py_Ref (*)(void*, int, int))array->f_get;
     res->f_set = (bool (*)(void*, int, int, py_Ref))array->f_set;
     return true;
 }
@@ -442,7 +523,7 @@ static bool _chunked_array2d_view(py_OutRef out,
     c11_array2d_view* res = _array2d_view__new(out, keepalive, start_col, start_row, width, height);
     if(res == NULL) return false;
     res->ctx = array;
-    res->f_get = (py_Ref(*)(void*, int, int))c11_chunked_array2d__get;
+    res->f_get = (py_Ref (*)(void*, int, int))c11_chunked_array2d__get;
     res->f_set = (bool (*)(void*, int, int, py_Ref))c11_chunked_array2d__set;
     return true;
 }
@@ -754,6 +835,7 @@ static void register_array2d_like(py_Ref mod) {
     py_bindmethod(type, "index", array2d_like_index);
 
     py_bindmethod(type, "render", array2d_like_render);
+    py_bindmethod(type, "render_with_color", array2d_like_render_with_color);
 
     py_bindmethod(type, "all", array2d_like_all);
     py_bindmethod(type, "any", array2d_like_any);

src/modules/picoterm.c

@@ -0,0 +1,98 @@
+#include "pocketpy/pocketpy.h"
+#include "pocketpy/objects/base.h"
+#include <stdio.h>
+#include "pocketpy/common/vector.h"
+
+static bool picoterm_enable_full_buffering_mode(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(0);
+    static char buf[1024 * 32];  // 32KB
+    setvbuf(stdout, buf, _IOFBF, sizeof(buf));
+    py_newnone(py_retval());
+    return true;
+}
+
+// https://gist.github.com/fnky/458719343aabd01cfb17a3a4f7296797
+typedef struct {
+    c11_sv text;
+    char suffix;
+} AnsiEscapedToken;
+
+static bool split_ansi_escaped_string(c11_sv sv, c11_vector* out_tokens);
+
+static bool picoterm_split_ansi_escaped_string(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    PY_CHECK_ARG_TYPE(0, tp_str);
+    c11_sv s = py_tosv(argv);
+    c11_vector /*T=AnsiEscapedToken*/ tokens;
+    c11_vector__ctor(&tokens, sizeof(AnsiEscapedToken));
+    if(!split_ansi_escaped_string(s, &tokens)) {
+        c11_vector__dtor(&tokens);
+        return ValueError("invalid ANSI escape sequences");
+    }
+    py_newlistn(py_retval(), tokens.length);
+    for(int i = 0; i < tokens.length; i++) {
+        AnsiEscapedToken t = c11__getitem(AnsiEscapedToken, &tokens, i);
+        py_ItemRef item = py_list_getitem(py_retval(), i);
+        py_newstrv(item, t.text);
+    }
+    c11_vector__dtor(&tokens);
+    return true;
+}
+
+void pk__add_module_picoterm() {
+    py_Ref mod = py_newmodule("picoterm");
+
+    py_bindfunc(mod, "enable_full_buffering_mode", picoterm_enable_full_buffering_mode);
+    py_bindfunc(mod, "split_ansi_escaped_string", picoterm_split_ansi_escaped_string);
+}
+
+static bool split_ansi_escaped_string(c11_sv sv, c11_vector* out_tokens) {
+    const char* p = sv.data;
+    int i = 0;
+    while(i < sv.size) {
+        if(p[i] == '\x1b') {
+            i++;  // skip '\x1b'
+            if(i >= sv.size || p[i] != '[') {
+                return false;  // invalid escape sequence
+            }
+
+            int esc_start = i - 1;
+            i++;  // skip '['
+
+            c11_sv content;
+            content.data = p + i;
+            while(i < sv.size && !((p[i] >= 'A' && p[i] <= 'Z') || (p[i] >= 'a' && p[i] <= 'z'))) {
+                i++;
+            }
+            content.size = p + i - content.data;
+            if(i >= sv.size) {
+                return false;  // invalid escape sequence
+            }
+
+            char suffix = p[i];
+            i++;  // skip suffix
+
+            AnsiEscapedToken token;
+            token.text = (c11_sv){p + esc_start, i - esc_start};
+            token.suffix = suffix;
+            c11_vector__push(AnsiEscapedToken, out_tokens, token);
+        } else if(p[i] == '\n') {
+            AnsiEscapedToken token;
+            token.text = (c11_sv){p + i, 1};
+            token.suffix = '\n';
+            c11_vector__push(AnsiEscapedToken, out_tokens, token);
+            i++;
+        } else {
+            int text_start = i;
+            while(i < sv.size && p[i] != '\x1b' && p[i] != '\n') {
+                i++;
+            }
+            AnsiEscapedToken token;
+            token.text = (c11_sv){p + text_start, i - text_start};
+            token.suffix = '\0';
+            c11_vector__push(AnsiEscapedToken, out_tokens, token);
+        }
+    }
+
+    return true;
+}

src/modules/pkpy.c

@@ -65,14 +65,6 @@ static bool pkpy_is_user_defined_type(int argc, py_Ref argv) {
     return true;
 }
 
-static bool pkpy_enable_full_buffering_mode(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(0);
-    static char buf[1024 * 128];
-    setvbuf(stdout, buf, _IOFBF, sizeof(buf));
-    py_newnone(py_retval());
-    return true;
-}
-
 static bool pkpy_currentvm(int argc, py_Ref argv) {
     PY_CHECK_ARGC(0);
     py_newint(py_retval(), py_currentvm());
@@ -539,7 +531,6 @@ void pk__add_module_pkpy() {
 
     py_bindfunc(mod, "memory_usage", pkpy_memory_usage);
     py_bindfunc(mod, "is_user_defined_type", pkpy_is_user_defined_type);
-    py_bindfunc(mod, "enable_full_buffering_mode", pkpy_enable_full_buffering_mode);
 
     py_bindfunc(mod, "currentvm", pkpy_currentvm);
 

src/modules/vmath.c

@@ -944,7 +944,7 @@ static bool color32_to_hex(int argc, py_Ref argv) {
     return true;
 }
 
-static void c11_color32_premult(c11_color32* color) {
+void c11_color32_premult(c11_color32* color) {
     if(color->a == 255) return;
     float alpha = color->a / 255.0f;
     color->r = (unsigned char)(color->r * alpha);

tests/92_picoterm.py

@@ -0,0 +1,23 @@
+import picoterm
+from vmath import rgb
+
+picoterm.enable_full_buffering_mode()
+
+bg = rgb(78, 118, 164)
+fg = rgb(200, 200, 0)
+text = "hello, \nworld"
+text = bg.ansi_bg(text)
+text = fg.ansi_fg(text)
+
+def ansi_italic(text: str):
+    return f'\x1b[3m{text}\x1b[0m'
+
+text = ansi_italic(text) + '123'
+print(text)
+
+cpnts = picoterm.split_ansi_escaped_string(text)
+
+assert cpnts == ['\x1b[3m', '\x1b[38;2;200;200;0m', '\x1b[48;2;78;118;164m', 'hello, ', '\n', 'world', '\x1b[0m', '\x1b[0m', '\x1b[0m', '123']
+
+cpnts_join = ''.join(cpnts)
+assert cpnts_join == text