add test for numpy F-contiguous constraint rejection (#847)

wjakob · Jan 8, 2025 · efbeb8a · efbeb8a
1 parent b95eb75
commit efbeb8a
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 0 deletions.
diff --git a/tests/test_ndarray.cpp b/tests/test_ndarray.cpp
@@ -351,6 +351,12 @@ NB_MODULE(test_ndarray_ext, m) {
 
     m.def("check", [](nb::handle h) { return nb::ndarray_check(h); });
 
+    m.def("accept_np_both_true_contig_a", 
+          [](nb::ndarray<float, nb::numpy, nb::shape<2, 1>, nb::any_contig> a) { return a(0, 0); });
+    m.def("accept_np_both_true_contig_c", 
+          [](nb::ndarray<float, nb::numpy, nb::shape<2, 1>, nb::c_contig> a) { return a(0, 0); });
+    m.def("accept_np_both_true_contig_f", 
+          [](nb::ndarray<float, nb::numpy, nb::shape<2, 1>, nb::f_contig> a) { return a(0, 0); });
 
     struct Cls {
         auto f1() { return nb::ndarray<nb::numpy, float>(data, { 10 }, nb::handle()); }

diff --git a/tests/test_ndarray.py b/tests/test_ndarray.py
@@ -957,3 +957,13 @@ def test51_return_from_stack():
     import numpy as np
     assert np.all(t.ret_from_stack_1() == [1,2,3])
     assert np.all(t.ret_from_stack_2() == [1,2,3])
+
+@needs_numpy
+def test52_accept_np_both_true_contig():
+    import numpy as np
+    a = np.zeros((2, 1), dtype=np.float32)
+    assert a.flags['C_CONTIGUOUS'] and a.flags['F_CONTIGUOUS']
+    t.accept_np_both_true_contig_a(a)
+    t.accept_np_both_true_contig_c(a)
+    t.accept_np_both_true_contig_f(a)
+
diff --git a/tests/test_ndarray_ext.pyi.ref b/tests/test_ndarray_ext.pyi.ref
@@ -27,6 +27,12 @@ class Matrix4f:
 
     def data_copy(self) -> Annotated[ArrayLike, dict(dtype='float32', shape=(4, 4), order='F')]: ...
 
+def accept_np_both_true_contig_a(arg: Annotated[ArrayLike, dict(dtype='float32', shape=(2, 1), order='A')], /) -> float: ...
+
+def accept_np_both_true_contig_c(arg: Annotated[ArrayLike, dict(dtype='float32', shape=(2, 1), order='C')], /) -> float: ...
+
+def accept_np_both_true_contig_f(arg: Annotated[ArrayLike, dict(dtype='float32', shape=(2, 1), order='F')], /) -> float: ...
+
 def accept_ro(arg: Annotated[ArrayLike, dict(dtype='float32', shape=(2), writable=False)], /) -> float: ...
 
 def accept_rw(arg: Annotated[ArrayLike, dict(dtype='float32', shape=(2))], /) -> float: ...