diff --git a/Tests/test_pyarrow.py b/Tests/test_pyarrow.py index f282f2c0059..2dfa440d181 100644 --- a/Tests/test_pyarrow.py +++ b/Tests/test_pyarrow.py @@ -84,21 +84,21 @@ def _test_img_equals_int32_pyarray( ).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]), - ), +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]), ) + + +@pytest.mark.parametrize("mode, dtype, mask", mode_dtype_mask) def test_to_array(mode: str, dtype: pyarrow.DataType, mask: list[int] | None) -> None: img = hopper(mode) @@ -218,6 +218,28 @@ def test_fromarray(mode: str, data_tp: DataShape, mask: list[int] | None) -> Non _test_img_equals_pyarray(img, arr, mask, elts_per_pixel) +@pytest.mark.parametrize("mode, dtype, mask", mode_dtype_mask) +def test_fromarray_to_array( + mode: str, dtype: pyarrow.DataType, mask: list[int] | None +) -> None: + img = hopper(mode) + + borrowed = Image.fromarrow(pyarrow.array(img), mode, img.size) # type: ignore[call-overload] + + arr = pyarrow.array(borrowed) # type: ignore[call-overload] + _test_img_equals_pyarray(img, arr, mask) + assert arr.type == dtype + + +def test_fromarray_interleaved_to_array() -> None: + img = hopper("RGBA") + + arr = pyarrow.array(list(img.tobytes()), type=pyarrow.uint8()) + borrowed = Image.fromarrow(arr, "RGBA", img.size) + + _test_img_equals_pyarray(img, pyarrow.array(borrowed), None) # type: ignore[call-overload] + + @pytest.mark.parametrize( "mode, data_tp, mask", ( diff --git a/src/libImaging/Arrow.c b/src/libImaging/Arrow.c index 3ca227d4f8e..b061dae4099 100644 --- a/src/libImaging/Arrow.c +++ b/src/libImaging/Arrow.c @@ -324,6 +324,8 @@ export_single_channel_array(Imaging im, struct ArrowArray *array) { if (im->block) { array->buffers[1] = im->block; + } else if (im->arrow_array_capsule) { + array->buffers[1] = im->image[0]; } else { array->buffers[1] = im->blocks[0].ptr; } @@ -407,6 +409,8 @@ export_fixed_pixel_array(Imaging im, struct ArrowArray *array) { if (im->block) { array->children[0]->buffers[1] = im->block; + } else if (im->arrow_array_capsule) { + array->children[0]->buffers[1] = im->image[0]; } else { array->children[0]->buffers[1] = im->blocks[0].ptr; }