Add test that returns nb::ndarray<nb::memview> (#1181)

Author: hpkfft

tests/test_ndarray.cpp

@@ -240,7 +240,8 @@ NB_MODULE(test_ndarray_ext, m) {
     });
 
     m.def("destruct_count", []() { return destruct_count; });
-    m.def("return_dlpack", []() {
+
+    m.def("return_no_framework", []() {
         float *f = new float[8] { 1, 2, 3, 4, 5, 6, 7, 8 };
         size_t shape[2] = { 2, 4 };
 
@@ -299,16 +300,28 @@ NB_MODULE(test_ndarray_ext, m) {
                                                                 deleter);
     });
 
+    m.def("ret_memview", []() {
+        double *d = new double[8] { 1, 2, 3, 4, 5, 6, 7, 8 };
+        size_t shape[2] = { 2, 4 };
+
+        nb::capsule deleter(d, [](void *data) noexcept {
+           destruct_count++;
+           delete[] (double *) data;
+        });
+
+        return nb::ndarray<nb::memview, double, nb::shape<2, 4>>(d, 2, shape,
+                                                                 deleter);
+    });
+
     m.def("ret_array_scalar", []() {
-            float* f = new float[1] { 1 };
-            size_t shape[1] = {};
+            float* f = new float{ 1.0f };
 
             nb::capsule deleter(f, [](void* data) noexcept {
                 destruct_count++;
-                delete[] (float *) data;
+                delete (float *) data;
             });
 
-            return nb::ndarray<nb::numpy, float>(f, 0, shape, deleter);
+            return nb::ndarray<nb::numpy, float>(f, 0, nullptr, deleter);
     });
 
     m.def("noop_3d_c_contig",

tests/test_ndarray.py

@@ -2,7 +2,7 @@
 import pytest
 import warnings
 import importlib
-from common import collect
+from common import collect, skip_on_pypy
 
 try:
     import numpy as np
@@ -209,50 +209,77 @@ def test11_implicit_conversion_pytorch():
         t.noimplicit(torch.zeros(2, 2, 10, dtype=torch.float32)[:, :, 4])
 
 
-def test14_destroy_capsule():
+@needs_numpy
+def test12_process_image():
+    x = np.arange(120, dtype=np.ubyte).reshape(8, 5, 3)
+    t.process(x)
+    assert np.all(x == np.arange(0, 240, 2, dtype=np.ubyte).reshape(8, 5, 3))
+
+
+def test13_destroy_capsule():
     collect()
     dc = t.destruct_count()
-    a = t.return_dlpack()
-    assert dc == t.destruct_count()
-    del a
+    capsule = t.return_no_framework()
+    assert 'dltensor' in repr(capsule)
+    assert 'versioned' not in repr(capsule)
+    assert t.destruct_count() == dc
+    del capsule
     collect()
     assert t.destruct_count() - dc == 1
 
 
 @needs_numpy
-def test15_consume_numpy():
+def test14_consume_numpy():
     collect()
     class wrapper:
         def __init__(self, value):
             self.value = value
         def __dlpack__(self):
             return self.value
     dc = t.destruct_count()
-    a = t.return_dlpack()
+    capsule = t.return_no_framework()
     if hasattr(np, '_from_dlpack'):
-        x = np._from_dlpack(wrapper(a))
+        x = np._from_dlpack(wrapper(capsule))
     elif hasattr(np, 'from_dlpack'):
-        x = np.from_dlpack(wrapper(a))
+        x = np.from_dlpack(wrapper(capsule))
     else:
         pytest.skip('your version of numpy is too old')
 
-    del a
+    del capsule
     collect()
     assert x.shape == (2, 4)
     assert np.all(x == [[1, 2, 3, 4], [5, 6, 7, 8]])
-    assert dc == t.destruct_count()
+    assert t.destruct_count() == dc
     del x
     collect()
     assert t.destruct_count() - dc == 1
 
 
 @needs_numpy
-def test16_passthrough():
+def test15_passthrough_numpy():
     a = t.ret_numpy()
     b = t.passthrough(a)
     assert a is b
 
-    a = np.array([1,2,3])
+    a = np.array([1, 2, 3])
+    b = t.passthrough(a)
+    assert a is b
+
+    a = None
+    with pytest.raises(TypeError) as excinfo:
+        b = t.passthrough(a)
+    assert 'incompatible function arguments' in str(excinfo.value)
+    b = t.passthrough_arg_none(a)
+    assert a is b
+
+
+@needs_torch
+def test16_passthrough_torch():
+    a = t.ret_pytorch()
+    b = t.passthrough(a)
+    assert a is b
+
+    a = torch.tensor([1, 2, 3])
     b = t.passthrough(a)
     assert a is b
 
@@ -292,6 +319,22 @@ def test18_return_pytorch():
     assert t.destruct_count() - dc == 1
 
 
+@skip_on_pypy
+def test19_return_memview():
+    collect()
+    dc = t.destruct_count()
+    x = t.ret_memview()
+    assert isinstance(x, memoryview)
+    assert x.itemsize == 8
+    assert x.ndim == 2
+    assert x.shape == (2, 4)
+    assert x.strides == (32, 8)  # in bytes
+    assert x.tolist() == [[1, 2, 3, 4], [5, 6, 7, 8]]
+    del x
+    collect()
+    assert t.destruct_count() - dc == 1
+
+
 @needs_numpy
 def test21_return_array_scalar():
     collect()
@@ -510,6 +553,7 @@ def test33_force_contig_numpy():
     assert b is not a
     assert np.all(b == a)
 
+
 @needs_torch
 @pytest.mark.filterwarnings
 def test34_force_contig_pytorch():
@@ -567,6 +611,7 @@ def test36_half():
     assert x.shape == (2, 4)
     assert np.all(x == [[1, 2, 3, 4], [5, 6, 7, 8]])
 
+
 @needs_numpy
 def test37_cast():
     a = t.cast(False)

tests/test_ndarray_ext.pyi.ref

@@ -97,7 +97,7 @@ def process(arg: Annotated[NDArray[numpy.uint8], dict(shape=(None, None, 3), ord
 
 def destruct_count() -> int: ...
 
-def return_dlpack() -> Annotated[NDArray[numpy.float32], dict(shape=(2, 4))]: ...
+def return_no_framework() -> Annotated[NDArray[numpy.float32], dict(shape=(2, 4))]: ...
 
 def passthrough(arg: NDArray, /) -> NDArray: ...
 
@@ -115,6 +115,8 @@ def ret_numpy_const() -> Annotated[NDArray[numpy.float32], dict(shape=(2, 4), wr
 
 def ret_pytorch() -> Annotated[NDArray[numpy.float32], dict(shape=(2, 4))]: ...
 
+def ret_memview() -> memoryview[dtype=float64, shape=(2, 4)]: ...
+
 def ret_array_scalar() -> NDArray[numpy.float32]: ...
 
 def noop_3d_c_contig(arg: Annotated[NDArray[numpy.float32], dict(shape=(None, None, None), order='C')], /) -> None: ...