WIP -- First light of round trip of image -> arrow -> image

* basic safety included
* only respecting the types we emit

Author: wiredfool

Tests/test_arrow.py

@@ -6,7 +6,7 @@
 
 from PIL import Image
 
-from .helper import assert_deep_equal, assert_image, hopper, skip_unless_feature
+from .helper import assert_deep_equal, assert_image, hopper, skip_unless_feature, assert_image_equal
 
 from typing import Any # undone
 
@@ -64,6 +64,11 @@ def test_to_array(mode: str, dtype: Any, mask: Any ) -> None:
     _test_img_equals_pyarray(img, arr, mask)
     assert arr.type == dtype
 
+    reloaded = Image.fromarrow(arr, mode, img.size)
+
+    assert reloaded
+
+    assert_image_equal(img, reloaded)
 
 def test_lifetime():
     # valgrind shouldn't error out here.

src/PIL/Image.py

@@ -3258,6 +3258,14 @@ class SupportsArrayInterface(Protocol):
     def __array_interface__(self) -> dict[str, Any]:
         raise NotImplementedError()
 
+class SupportsArrowArrayInterface(Protocol):
+    """
+    An object that has an ``__arrow_c_array__`` method corresponding to the arrow c data interface.
+    """
+
+    def __arrow_c_array__(self, requested_schema:"PyCapsule"=None) -> tuple["PyCapsule", "PyCapsule"]:
+        raise NotImplementedError()
+
 
 def fromarray(obj: SupportsArrayInterface, mode: str | None = None) -> Image:
     """
@@ -3347,6 +3355,18 @@ def fromarray(obj: SupportsArrayInterface, mode: str | None = None) -> Image:
     return frombuffer(mode, size, obj, "raw", rawmode, 0, 1)
 
 
+def fromarrow(obj: SupportsArrowArrayIngerface, mode, size) -> ImageFile.ImageFile:
+    if not hasattr(obj, '__arrow_c_array__'):
+        raise ValueError("arrow_c_array interface not found")
+
+    (schema_capsule, array_capsule) = obj.__arrow_c_array__()
+    _im = core.new_arrow(mode, size, schema_capsule, array_capsule)
+    if (_im):
+        return Image()._new(_im)
+
+    return None
+
+
 def fromqimage(im: ImageQt.QImage) -> ImageFile.ImageFile:
     """Creates an image instance from a QImage image"""
     from . import ImageQt

src/_imaging.c

@@ -89,7 +89,6 @@
 #endif
 
 #include "libImaging/Imaging.h"
-#include "libImaging/Arrow.h"
 
 #define _USE_MATH_DEFINES
 #include <math.h>
@@ -261,6 +260,38 @@ PyObject* ExportArrowArrayPyCapsule(ImagingObject *self) {
 }
 
 
+static PyObject *
+_new_arrow(PyObject *self, PyObject *args) {
+    char *mode;
+    int xsize, ysize;
+    PyObject *schema_capsule, *array_capsule;
+    PyObject *ret;
+
+    if (!PyArg_ParseTuple(args, "s(ii)OO", &mode, &xsize, &ysize,
+                          &schema_capsule, &array_capsule)) {
+        return NULL;
+    }
+
+    struct ArrowSchema* schema =
+      (struct ArrowSchema*)PyCapsule_GetPointer(schema_capsule, "arrow_schema");
+
+    struct ArrowArray* array =
+      (struct ArrowArray*)PyCapsule_GetPointer(array_capsule, "arrow_array");
+
+    ret = PyImagingNew(ImagingNewArrow(mode, xsize, ysize, schema, array));
+    if (schema->release){
+      schema->release(schema);
+      schema->release = NULL;
+    }
+    if (!ret && array->release) {
+      array->release(array);
+      array->release = NULL;
+    }
+    return ret;
+}
+
+
+
 /* -------------------------------------------------------------------- */
 /* EXCEPTION REROUTING                                                  */
 /* -------------------------------------------------------------------- */
@@ -4258,6 +4289,7 @@ static PyMethodDef functions[] = {
     {"fill", (PyCFunction)_fill, METH_VARARGS},
     {"new", (PyCFunction)_new, METH_VARARGS},
     {"new_block", (PyCFunction)_new_block, METH_VARARGS},
+    {"new_arrow", (PyCFunction)_new_arrow, METH_VARARGS},
     {"merge", (PyCFunction)_merge, METH_VARARGS},
 
     /* Functions */

src/libImaging/Imaging.h

@@ -20,6 +20,8 @@ extern "C" {
 #define M_PI 3.1415926535897932384626433832795
 #endif
 
+#include "Arrow.h"
+
 /* -------------------------------------------------------------------- */
 
 /*
@@ -107,11 +109,17 @@ struct ImagingMemoryInstance {
 
     /* arrow */
     int arrow_borrow; /* Number of arrow arrays that have been allocated */
+    char band_names[4][3]; /* names of bands, max 2 char + null terminator */
+    char arrow_band_format[2]; /* single character + null terminator */
+
+    int read_only;  /* flag for read-only. set for arrow borrowed arrays */
+    struct ArrowArray *arrow_array_capsule; /* upstream arrow array source */
+
     int blocks_count; /* Number of blocks that have been allocated */
     int lines_per_block; /* Number of lines in a block have been allocated */
 
-    char band_names[4][3]; /* names of bands, max 2 char + null terminator */
-    char arrow_band_format[2]; /* single character + null terminator */
+
+
 };
 
 #define IMAGING_PIXEL_1(im, x, y) ((im)->image8[(y)][(x)])
@@ -195,6 +203,11 @@ ImagingDelete(Imaging im);
 extern Imaging
 ImagingNewBlock(const char *mode, int xsize, int ysize);
 
+extern Imaging
+ImagingNewArrow(const char *mode, int xsize, int ysize,
+                struct ArrowSchema *schema,
+                struct ArrowArray *external_array);
+
 extern Imaging
 ImagingNewPrologue(const char *mode, int xsize, int ysize);
 extern Imaging
@@ -712,7 +725,6 @@ _imaging_tell_pyFd(PyObject *fd);
 
 /* Arrow */
 
-#include "Arrow.h"
 extern int export_imaging_array(Imaging im, struct ArrowArray* array);
 extern int export_imaging_schema(Imaging im, struct ArrowSchema* schema);
 extern void export_uint32_type(struct ArrowSchema* schema);

src/libImaging/Storage.c

@@ -278,6 +278,7 @@ ImagingNewPrologueSubtype(const char *mode, int xsize, int ysize, int size) {
             break;
     }
 
+    // UNDONE -- not accurate for arrow
     MUTEX_LOCK(&ImagingDefaultArena.mutex);
     ImagingDefaultArena.stats_new_count += 1;
     MUTEX_UNLOCK(&ImagingDefaultArena.mutex);
@@ -556,6 +557,62 @@ ImagingAllocateBlock(Imaging im) {
     return im;
 }
 
+/* Borrowed Arrow Storage Type */
+/* --------------------------- */
+/* Don't allocate the image. */
+
+
+static void
+ImagingDestroyArrow(Imaging im) {
+    if (im->arrow_array_capsule && im->arrow_array_capsule->release) {
+        im->arrow_array_capsule->release(im->arrow_array_capsule);
+        im->arrow_array_capsule->release = NULL;
+        im->arrow_array_capsule = NULL;
+    }
+}
+
+Imaging
+ImagingAllocateArrow(Imaging im, struct ArrowArray *external_array) {
+    /* don't really allocate, but naming patterns */
+    Py_ssize_t y, i;
+
+    char* borrowed_buffer = NULL;
+    struct ArrowArray* arr = external_array;
+
+    /* overflow check for malloc */
+    if (im->linesize && im->ysize > INT_MAX / im->linesize) {
+        return (Imaging)ImagingError_MemoryError();
+    }
+
+    if (arr->n_children == 1) {
+        arr = arr->children[0];
+    }
+    if (arr->n_buffers == 2) {
+        // undone offset. offset here is # of elements,
+        // so depends on the size of the element
+        // undone image is char*, arrow is void *
+        // buffer 0 is the null list
+        // buffer 1 is the data
+        borrowed_buffer = (char *)arr->buffers[1];
+    }
+
+    if (! borrowed_buffer) {
+        // UNDONE better error here.
+        // currently, only wanting one where
+        return (Imaging)ImagingError_MemoryError();
+    }
+
+    for (y = i = 0; y < im->ysize; y++) {
+        im->image[y] = borrowed_buffer + i;
+        i += im->linesize;
+    }
+    im->read_only = 1;
+    im->arrow_array_capsule = external_array;
+    im->destroy = ImagingDestroyArrow;
+
+    return im;
+}
+
 /* --------------------------------------------------------------------
  * Create a new, internally allocated, image.
  */
@@ -627,6 +684,53 @@ ImagingNewBlock(const char *mode, int xsize, int ysize) {
     return NULL;
 }
 
+Imaging
+ImagingNewArrow(const char *mode, int xsize, int ysize,
+                struct ArrowSchema *schema,
+                struct ArrowArray *external_array) {
+    /* A borrowed arrow array */
+    Imaging im;
+
+    if (xsize < 0 || ysize < 0) {
+        return (Imaging)ImagingError_ValueError("bad image size");
+    }
+
+    im = ImagingNewPrologue(mode, xsize, ysize);
+    if (!im) {
+        return NULL;
+    }
+
+    int64_t pixels = (int64_t)xsize * (int64_t)ysize;
+
+    if (((strcmp(schema->format, "i") == 0 &&
+          im->pixelsize == 4) ||
+         (strcmp(schema->format, im->arrow_band_format) == 0 &&
+          im->bands == 1))
+        && pixels == external_array->length) {
+        // one arrow element per, and it matches a pixelsize*char
+        if (ImagingAllocateArrow(im, external_array)) {
+            return im;
+        }
+    }
+    if (strcmp(schema->format, "+w:4") == 0
+        && im->pixelsize == 4
+        && schema->n_children > 0
+        && schema->children
+        && strcmp(schema->children[0]->format, "C") == 0
+        && pixels == external_array->length
+        && external_array->n_children == 1
+        && external_array->children
+        && 4 * pixels == external_array->children[0]->length) {
+        // 4 up element of char into pixelsize == 4
+        if (ImagingAllocateArrow(im, external_array)) {
+            return im;
+        }
+    }
+
+    ImagingDelete(im);
+    return NULL;
+}
+
 Imaging
 ImagingNew2Dirty(const char *mode, Imaging imOut, Imaging imIn) {
     /* allocate or validate output image */