Image -> Arrow support

Tests/test_arrow.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import Any  # undone
+
+import pytest
+
+from PIL import Image
+
+from .helper import (
+    assert_deep_equal,
+    assert_image_equal,
+    hopper,
+)
+
+pyarrow = pytest.importorskip("pyarrow", reason="PyArrow not installed")
+
+TEST_IMAGE_SIZE = (10, 10)
+
+
+def _test_img_equals_pyarray(img: Image.Image, arr: Any, mask) -> None:
+    assert img.height * img.width == len(arr)
+    px = img.load()
+    assert px is not None
+    for x in range(0, img.size[0], int(img.size[0] / 10)):
+        for y in range(0, img.size[1], int(img.size[1] / 10)):
+            if mask:
+                for ix, elt in enumerate(mask):
+                    assert px[x, y][ix] == arr[y * img.width + x].as_py()[elt]
+            else:
+                assert_deep_equal(px[x, y], arr[y * img.width + x].as_py())
+
+
+# really hard to get a non-nullable list type
+fl_uint8_4_type = pyarrow.field(
+    "_", pyarrow.list_(pyarrow.field("_", pyarrow.uint8()).with_nullable(False), 4)
+).type
+
+
+@pytest.mark.parametrize(
+    "mode, dtype, mask",
+    (
+        ("L", pyarrow.uint8(), None),
+        ("I", pyarrow.int32(), None),
+        ("F", pyarrow.float32(), None),
+        ("LA", fl_uint8_4_type, [0, 3]),
+        ("RGB", fl_uint8_4_type, [0, 1, 2]),
+        ("RGBA", fl_uint8_4_type, None),
+        ("RGBX", fl_uint8_4_type, None),
+        ("CMYK", fl_uint8_4_type, None),
+        ("YCbCr", fl_uint8_4_type, [0, 1, 2]),
+        ("HSV", fl_uint8_4_type, [0, 1, 2]),
+    ),
+)
+def test_to_array(mode: str, dtype: Any, mask: Any) -> None:
+    img = hopper(mode)
+
+    # Resize to non-square
+    img = img.crop((3, 0, 124, 127))
+    assert img.size == (121, 127)
+
+    arr = pyarrow.array(img)
+    _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.
+    # arrays should be accessible after the image is deleted.
+
+    img = hopper("L")
+
+    arr_1 = pyarrow.array(img)
+    arr_2 = pyarrow.array(img)
+
+    del img
+
+    assert arr_1.sum().as_py() > 0
+    del arr_1
+
+    assert arr_2.sum().as_py() > 0
+    del arr_2
+
+
+def test_lifetime2():
+    # valgrind shouldn't error out here.
+    # img should remain after the arrays are collected.
+
+    img = hopper("L")
+
+    arr_1 = pyarrow.array(img)
+    arr_2 = pyarrow.array(img)
+
+    assert arr_1.sum().as_py() > 0
+    del arr_1
+
+    assert arr_2.sum().as_py() > 0
+    del arr_2
+
+    img2 = img.copy()
+    px = img2.load()
+    assert isinstance(px[0, 0], int)

pyproject.toml

@@ -61,6 +61,7 @@ optional-dependencies.tests = [
   "markdown2",
   "olefile",
   "packaging",
+  "pyarrow",
   "pyroma",
   "pytest",
   "pytest-cov",

setup.py

@@ -64,6 +64,7 @@ def get_version() -> str:
 _LIB_IMAGING = (
     "Access",
     "AlphaComposite",
+    "Arrow",
     "Resample",
     "Reduce",
     "Bands",

src/PIL/Image.py

@@ -748,6 +748,16 @@ def __array_interface__(self) -> dict[str, str | bytes | int | tuple[int, ...]]:
         new["shape"], new["typestr"] = _conv_type_shape(self)
         return new
 
+    def __arrow_c_schema__(self) -> object:
+        self.load()
+        return self.im.__arrow_c_schema__()
+
+    def __arrow_c_array__(
+        self, requested_schema: object | None = None
+    ) -> Tuple[object, object]:
+        self.load()
+        return (self.im.__arrow_c_schema__(), self.im.__arrow_c_array__())
+
     def __getstate__(self) -> list[Any]:
         im_data = self.tobytes()  # load image first
         return [self.info, self.mode, self.size, self.getpalette(), im_data]
@@ -3250,6 +3260,17 @@ 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:
     """
     Creates an image memory from an object exporting the array interface
@@ -3338,6 +3359,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

@@ -223,6 +223,77 @@
     return PyObject_GetBuffer(buffer, view, PyBUF_SIMPLE);
 }
 
+/* -------------------------------------------------------------------- */
+/* Arrow HANDLING                                                       */
+/* -------------------------------------------------------------------- */
+
+void
+ReleaseArrowSchemaPyCapsule(PyObject *capsule) {
+    struct ArrowSchema *schema =
+        (struct ArrowSchema *)PyCapsule_GetPointer(capsule, "arrow_schema");
+    if (schema->release != NULL) {
+        schema->release(schema);
+    }
+    free(schema);
+}
+
+PyObject *
+ExportArrowSchemaPyCapsule(ImagingObject *self) {
+    struct ArrowSchema *schema =
+        (struct ArrowSchema *)calloc(1, sizeof(struct ArrowSchema));
+    export_imaging_schema(self->image, schema);
+    return PyCapsule_New(schema, "arrow_schema", ReleaseArrowSchemaPyCapsule);
+}
+
+void
+ReleaseArrowArrayPyCapsule(PyObject *capsule) {
+    struct ArrowArray *array =
+        (struct ArrowArray *)PyCapsule_GetPointer(capsule, "arrow_array");
+    if (array->release != NULL) {
+        array->release(array);
+    }
+    free(array);
+}
+
+PyObject *
+ExportArrowArrayPyCapsule(ImagingObject *self) {
+    struct ArrowArray *array =
+        (struct ArrowArray *)calloc(1, sizeof(struct ArrowArray));
+    export_imaging_array(self->image, array);
+    return PyCapsule_New(array, "arrow_array", ReleaseArrowArrayPyCapsule);
+}
+
+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                                                  */
 /* -------------------------------------------------------------------- */
@@ -3678,6 +3749,10 @@
     /* Misc. */
     {"save_ppm", (PyCFunction)_save_ppm, METH_VARARGS},
 
+    /* arrow */
+    {"__arrow_c_schema__", (PyCFunction)ExportArrowSchemaPyCapsule, METH_VARARGS},
+    {"__arrow_c_array__", (PyCFunction)ExportArrowArrayPyCapsule, METH_VARARGS},
+
     {NULL, NULL} /* sentinel */
 };
 
@@ -4216,6 +4291,7 @@
     {"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 */