test_numpy_dtypes: Add test for py::vectorize()

Author: EricCousineau-TRI

tests/test_numpy_dtypes.cpp

@@ -462,10 +462,16 @@ TEST_SUBMODULE(numpy_dtypes, m) {
     m.def("buffer_to_dtype", [](py::buffer& buf) { return py::dtype(buf.request()); });
 
     // test_scalar_conversion
-    m.def("f_simple", [](SimpleStruct s) { return s.uint_ * 10; });
+    auto f_simple = [](SimpleStruct s) { return s.uint_ * 10; };
+    m.def("f_simple", f_simple);
     m.def("f_packed", [](PackedStruct s) { return s.uint_ * 10; });
     m.def("f_nested", [](NestedStruct s) { return s.a.uint_ * 10; });
 
+    // test_vectorize
+    m.def("f_simple_vectorized", py::vectorize(f_simple));
+    auto f_simple_pass_thru = [](SimpleStruct s) { return s; };
+    m.def("f_simple_pass_thru_vectorized", py::vectorize(f_simple_pass_thru));
+
     // test_register_dtype
     m.def("register_dtype", []() { PYBIND11_NUMPY_DTYPE(SimpleStruct, bool_, uint_, float_, ldbl_); });
 

tests/test_numpy_dtypes.py

@@ -287,6 +287,15 @@ def test_scalar_conversion():
                 assert 'incompatible function arguments' in str(excinfo.value)
 
 
+def test_vectorize():
+    n = 3
+    array = m.create_rec_simple(n)
+    values = m.f_simple_vectorized(array)
+    np.testing.assert_array_equal(values, [0, 10, 20])
+    array_2 = m.f_simple_pass_thru_vectorized(array)
+    np.testing.assert_array_equal(array, array_2)
+
+
 def test_register_dtype():
     with pytest.raises(RuntimeError) as excinfo:
         m.register_dtype()