Update cute_png.c

3rd/cute_png/src/cute_png.c

@@ -23,6 +23,7 @@ static bool cute_png_loads(int argc, py_Ref argv) {
             py_newcolor32(slot, color);
         }
     }
+    cp_free_png(&image);
     return true;
 }
 
@@ -45,12 +46,156 @@ static bool cute_png_dumps(int argc, py_Ref argv) {
     assert(saved_image.data != NULL);
     unsigned char* data = py_newbytes(py_retval(), saved_image.size);
     memcpy(data, saved_image.data, saved_image.size);
+    CUTE_PNG_FREE(saved_image.data);
+    cp_free_png(&image);
+    return true;
+}
+
+static void cute_png_Image__dtor(void* ud) {
+    cp_image_t* image = (cp_image_t*)ud;
+    cp_free_png(image);
+}
+
+static bool cute_png_Image__new__(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(3);
+    PY_CHECK_ARG_TYPE(1, tp_int);
+    PY_CHECK_ARG_TYPE(2, tp_int);
+    int width = py_toint(py_arg(1));
+    int height = py_toint(py_arg(2));
+    cp_image_t* ud = py_newobject(py_retval(), py_totype(argv), 0, sizeof(cp_image_t));
+    *ud = cp_load_blank(width, height);
+    return true;
+}
+
+static bool cute_png_Image__from_bytes_STATIC(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    PY_CHECK_ARG_TYPE(0, tp_bytes);
+    int size;
+    unsigned char* data = py_tobytes(argv, &size);
+    cp_image_t image = cp_load_png_mem(data, size);
+    if(image.pix == NULL) return ValueError("cute_png: %s", cp_error_reason);
+    cp_image_t* ud =
+        py_newobject(py_retval(), py_gettype("cute_png", py_name("Image")), 0, sizeof(cp_image_t));
+    *ud = image;
+    return true;
+}
+
+static bool cute_png_Image__width(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    cp_image_t* image = py_touserdata(argv);
+    py_newint(py_retval(), image->w);
+    return true;
+}
+
+static bool cute_png_Image__height(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    cp_image_t* image = py_touserdata(argv);
+    py_newint(py_retval(), image->h);
+    return true;
+}
+
+static bool cute_png_Image__setpixel(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(4);
+    cp_image_t* image = py_touserdata(argv);
+    PY_CHECK_ARG_TYPE(1, tp_int);
+    PY_CHECK_ARG_TYPE(2, tp_int);
+    PY_CHECK_ARG_TYPE(3, tp_color32);
+    int x = py_toint(py_arg(1));
+    int y = py_toint(py_arg(2));
+    c11_color32 color = py_tocolor32(py_arg(3));
+    if(x < 0 || x >= image->w || y < 0 || y >= image->h) {
+        return IndexError("cute_png.Image: index out of range");
+    }
+    cp_pixel_t pixel = cp_make_pixel_a(color.r, color.g, color.b, color.a);
+    image->pix[y * image->w + x] = pixel;
+    py_newnone(py_retval());
+    return true;
+}
+
+static bool cute_png_Image__getpixel(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(3);
+    cp_image_t* image = py_touserdata(argv);
+    PY_CHECK_ARG_TYPE(1, tp_int);
+    PY_CHECK_ARG_TYPE(2, tp_int);
+    int x = py_toint(py_arg(1));
+    int y = py_toint(py_arg(2));
+    if(x < 0 || x >= image->w || y < 0 || y >= image->h) {
+        return IndexError("cute_png.Image: index out of range");
+    }
+    cp_pixel_t pixel = image->pix[y * image->w + x];
+    c11_color32 color;
+    color.r = pixel.r;
+    color.g = pixel.g;
+    color.b = pixel.b;
+    color.a = pixel.a;
+    py_newcolor32(py_retval(), color);
+    return true;
+}
+
+static bool cute_png_Image__clear(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(2);
+    cp_image_t* image = py_touserdata(argv);
+    PY_CHECK_ARG_TYPE(1, tp_color32);
+    c11_color32 color = py_tocolor32(py_arg(1));
+    cp_pixel_t pixel = cp_make_pixel_a(color.r, color.g, color.b, color.a);
+    for(int y = 0; y < image->h; y++) {
+        for(int x = 0; x < image->w; x++) {
+            image->pix[y * image->w + x] = pixel;
+        }
+    }
+    py_newnone(py_retval());
+    return true;
+}
+
+static bool cute_png_Image__to_rgb565_bytes(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    cp_image_t* image = py_touserdata(argv);
+    unsigned char* data = py_newbytes(py_retval(), image->w * image->h * 2);
+    for(int y = 0; y < image->h; y++) {
+        for(int x = 0; x < image->w; x++) {
+            size_t idx = y * image->w + x;
+            cp_pixel_t pixel = image->pix[idx];
+            uint16_t r = (pixel.r >> 3) & 0x1F;
+            uint16_t g = (pixel.g >> 2) & 0x3F;
+            uint16_t b = (pixel.b >> 3) & 0x1F;
+            uint16_t rgb565 = (r << 11) | (g << 5) | b;
+            // use little-endian
+            data[idx * 2 + 0] = rgb565 & 0xFF;
+            data[idx * 2 + 1] = (rgb565 >> 8) & 0xFF;
+        }
+    }
+    return true;
+}
+
+static bool cute_png_Image__to_png_bytes(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    cp_image_t* image = py_touserdata(argv);
+    cp_saved_png_t saved_image = cp_save_png_to_memory(image);
+    assert(saved_image.data != NULL);
+    unsigned char* data = py_newbytes(py_retval(), saved_image.size);
+    memcpy(data, saved_image.data, saved_image.size);
+    CUTE_PNG_FREE(saved_image.data);
     return true;
 }
 
 void pk__add_module_cute_png() {
     py_GlobalRef mod = py_newmodule("cute_png");
+    py_Type tp_image = py_newtype("Image", tp_object, mod, cute_png_Image__dtor);
+    py_tpsetfinal(tp_image);
 
     py_bindfunc(mod, "loads", cute_png_loads);
     py_bindfunc(mod, "dumps", cute_png_dumps);
+
+    py_bindmethod(tp_image, "__new__", cute_png_Image__new__);
+    py_bindstaticmethod(tp_image, "from_bytes", cute_png_Image__from_bytes_STATIC);
+
+    py_bindproperty(tp_image, "width", cute_png_Image__width, NULL);
+    py_bindproperty(tp_image, "height", cute_png_Image__height, NULL);
+
+    py_bindmethod(tp_image, "setpixel", cute_png_Image__setpixel);
+    py_bindmethod(tp_image, "getpixel", cute_png_Image__getpixel);
+    py_bindmethod(tp_image, "clear", cute_png_Image__clear);
+
+    py_bindmethod(tp_image, "to_rgb565_bytes", cute_png_Image__to_rgb565_bytes);
+    py_bindmethod(tp_image, "to_png_bytes", cute_png_Image__to_png_bytes);
 }

include/typings/cute_png.pyi

@@ -1,5 +1,23 @@
 from array2d import array2d
 from vmath import color32
 
+class Image:
+    @property
+    def width(self) -> int: ...
+    @property
+    def height(self) -> int: ...
+
+    def __new__(cls, width: int, height: int) -> "Image": ...
+
+    @staticmethod
+    def from_bytes(data: bytes) -> "Image": ...
+
+    def setpixel(self, x: int, y: int, color: color32) -> None: ...
+    def getpixel(self, x: int, y: int) -> color32: ...
+    def clear(self, color: color32) -> None: ...
+
+    def to_rgb565_bytes(self) -> bytes: ...
+    def to_png_bytes(self) -> bytes: ...
+
 def loads(data: bytes) -> array2d[color32]: ...
 def dumps(image: array2d[color32]) -> bytes: ...

include/typings/vmath.pyi

@@ -213,6 +213,8 @@ class color32:
     def to_vec3(self) -> vec3: ...
     def to_vec3i(self) -> vec3i: ...
 
+    def to_rgb565(self) -> int: ...
+
     def ansi_fg(self, text: str) -> str: ...
     def ansi_bg(self, text: str) -> str: ...
 

src/modules/vmath.c

@@ -976,6 +976,18 @@ static bool color32_to_vec3i(int argc, py_Ref argv) {
     return true;
 }
 
+static bool color32_to_rgb565(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    c11_color32 color = py_tocolor32(argv);
+    c11_color32_premult(&color);
+    uint16_t r = (color.r >> 3) & 0x1F;
+    uint16_t g = (color.g >> 2) & 0x3F;
+    uint16_t b = (color.b >> 3) & 0x1F;
+    uint16_t rgb565 = (r << 11) | (g << 5) | b;
+    py_newint(py_retval(), rgb565);
+    return true;
+}
+
 static bool color32__eq__(int argc, py_Ref argv) {
     PY_CHECK_ARGC(2);
     if(argv[1].type != tp_color32) {
@@ -1279,6 +1291,7 @@ void pk__add_module_vmath() {
     py_bindmethod(color32, "to_hex", color32_to_hex);
     py_bindmethod(color32, "to_vec3", color32_to_vec3);
     py_bindmethod(color32, "to_vec3i", color32_to_vec3i);
+    py_bindmethod(color32, "to_rgb565", color32_to_rgb565);
     py_bindmethod(color32, "ansi_fg", color32_ansi_fg);
     py_bindmethod(color32, "ansi_bg", color32_ansi_bg);
     py_bindfunc(mod, "rgb", vmath_rgb);