Added tests for ArrayLike in signatures and fixed wrong types for Refs

pybind · Feb 14, 2025 · 0b49140 · 0b49140
1 parent 6771709
commit 0b49140
Show file tree

Hide file tree

Showing 7 changed files with 82 additions and 9 deletions.
diff --git a/include/pybind11/eigen/matrix.h b/include/pybind11/eigen/matrix.h
@@ -441,7 +441,9 @@ struct eigen_map_caster {
         }
     }
 
-    static constexpr auto name = props::descriptor;
+    // return_descr forces the use of NDArray instead of ArrayLike in args
+    // since Ref<...> args can only accept arrays.
+    static constexpr auto name = return_descr(props::descriptor);
 
     // Explicitly delete these: support python -> C++ conversion on these (i.e. these can be return
     // types but not bound arguments).  We still provide them (with an explicitly delete) so that

diff --git a/include/pybind11/eigen/tensor.h b/include/pybind11/eigen/tensor.h
@@ -505,7 +505,10 @@ struct type_caster<Eigen::TensorMap<Type, Options>,
     std::unique_ptr<MapType> value;
 
 public:
-    static constexpr auto name = get_tensor_descriptor<Type, true, needs_writeable>::value;
+    // return_descr forces the use of NDArray instead of ArrayLike since refs can only reference
+    // arrays
+    static constexpr auto name
+        = return_descr(get_tensor_descriptor<Type, true, needs_writeable>::value);
     explicit operator MapType *() { return value.get(); }
     explicit operator MapType &() { return *value; }
     explicit operator MapType &&() && { return std::move(*value); }

diff --git a/tests/test_eigen_matrix.cpp b/tests/test_eigen_matrix.cpp
@@ -440,4 +440,8 @@ TEST_SUBMODULE(eigen_matrix, m) {
         py::module_::import("numpy").attr("ones")(10);
         return v[0](5);
     });
+    m.def("round_trip_vector", [](const Eigen::VectorXf &x) -> Eigen::VectorXf { return x; });
+    m.def("round_trip_dense", [](const DenseMatrixR &m) -> DenseMatrixR { return m; });
+    m.def("round_trip_dense_ref",
+          [](const Eigen::Ref<DenseMatrixR> &m) -> Eigen::Ref<DenseMatrixR> { return m; });
 }
diff --git a/tests/test_eigen_matrix.py b/tests/test_eigen_matrix.py
@@ -95,19 +95,19 @@ def test_mutator_descriptors():
     with pytest.raises(TypeError) as excinfo:
         m.fixed_mutator_r(zc)
     assert (
-        '(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[5, 6]",'
+        '(arg0: typing.Annotated[numpy.typing.NDArray[numpy.float32], "[5, 6]",'
         ' "flags.writeable", "flags.c_contiguous"]) -> None' in str(excinfo.value)
     )
     with pytest.raises(TypeError) as excinfo:
         m.fixed_mutator_c(zr)
     assert (
-        '(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[5, 6]",'
+        '(arg0: typing.Annotated[numpy.typing.NDArray[numpy.float32], "[5, 6]",'
         ' "flags.writeable", "flags.f_contiguous"]) -> None' in str(excinfo.value)
     )
     with pytest.raises(TypeError) as excinfo:
         m.fixed_mutator_a(np.array([[1, 2], [3, 4]], dtype="float32"))
     assert (
-        '(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[5, 6]", "flags.writeable"]) -> None'
+        '(arg0: typing.Annotated[numpy.typing.NDArray[numpy.float32], "[5, 6]", "flags.writeable"]) -> None'
         in str(excinfo.value)
     )
     zr.flags.writeable = False
@@ -202,7 +202,7 @@ def test_negative_stride_from_python(msg):
         msg(excinfo.value)
         == """
         double_threer(): incompatible function arguments. The following argument types are supported:
-            1. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[1, 3]", "flags.writeable"]) -> None
+            1. (arg0: typing.Annotated[numpy.typing.NDArray[numpy.float32], "[1, 3]", "flags.writeable"]) -> None
 
         Invoked with: """
         + repr(np.array([5.0, 4.0, 3.0], dtype="float32"))
@@ -214,7 +214,7 @@ def test_negative_stride_from_python(msg):
         msg(excinfo.value)
         == """
         double_threec(): incompatible function arguments. The following argument types are supported:
-            1. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[3, 1]", "flags.writeable"]) -> None
+            1. (arg0: typing.Annotated[numpy.typing.NDArray[numpy.float32], "[3, 1]", "flags.writeable"]) -> None
 
         Invoked with: """
         + repr(np.array([7.0, 4.0, 1.0], dtype="float32"))
@@ -818,3 +818,22 @@ def test_custom_operator_new():
     o = m.CustomOperatorNew()
     np.testing.assert_allclose(o.a, 0.0)
     np.testing.assert_allclose(o.b.diagonal(), 1.0)
+
+
+def test_arraylike_signature(doc):
+    assert doc(m.round_trip_vector) == (
+        'round_trip_vector(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[m, 1]"])'
+        ' -> typing.Annotated[numpy.typing.NDArray[numpy.float32], "[m, 1]"]'
+    )
+    assert doc(m.round_trip_dense) == (
+        'round_trip_dense(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[m, n]"])'
+        ' -> typing.Annotated[numpy.typing.NDArray[numpy.float32], "[m, n]"]'
+    )
+    assert doc(m.round_trip_dense_ref) == (
+        'round_trip_dense_ref(arg0: typing.Annotated[numpy.typing.NDArray[numpy.float32], "[m, n]", "flags.writeable", "flags.c_contiguous"])'
+        ' -> typing.Annotated[numpy.typing.NDArray[numpy.float32], "[m, n]", "flags.writeable", "flags.c_contiguous"]'
+    )
+    m.round_trip_vector([1.0, 2.0])
+    m.round_trip_dense([[1.0, 2.0], [3.0, 4.0]])
+    with pytest.raises(TypeError, match="incompatible function arguments"):
+        m.round_trip_dense_ref([[1.0, 2.0], [3.0, 4.0]])
diff --git a/tests/test_eigen_tensor.py b/tests/test_eigen_tensor.py
@@ -285,10 +285,32 @@ def test_doc_string(m, doc):
 
     order_flag = f'"flags.{m.needed_options.lower()}_contiguous"'
     assert doc(m.round_trip_view_tensor) == (
-        f'round_trip_view_tensor(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float64, "[?, ?, ?]", "flags.writeable", {order_flag}])'
+        f'round_trip_view_tensor(arg0: typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]", "flags.writeable", {order_flag}])'
         f' -> typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]", "flags.writeable", {order_flag}]'
     )
     assert doc(m.round_trip_const_view_tensor) == (
-        f'round_trip_const_view_tensor(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float64, "[?, ?, ?]", {order_flag}])'
+        f'round_trip_const_view_tensor(arg0: typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]", {order_flag}])'
         ' -> typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]"]'
     )
+
+
+@pytest.mark.parametrize("m", submodules)
+def test_arraylike_signature(m, doc):
+    order_flag = f'"flags.{m.needed_options.lower()}_contiguous"'
+    assert doc(m.round_trip_tensor) == (
+        'round_trip_tensor(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float64, "[?, ?, ?]"])'
+        ' -> typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]"]'
+    )
+    assert doc(m.round_trip_tensor_noconvert) == (
+        'round_trip_tensor_noconvert(tensor: typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]"])'
+        ' -> typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]"]'
+    )
+    assert doc(m.round_trip_view_tensor) == (
+        f'round_trip_view_tensor(arg0: typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]", "flags.writeable", {order_flag}])'
+        f' -> typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]", "flags.writeable", {order_flag}]'
+    )
+    m.round_trip_tensor(tensor_ref.tolist())
+    with pytest.raises(TypeError, match="incompatible function arguments"):
+        m.round_trip_tensor_noconvert(tensor_ref.tolist())
+    with pytest.raises(TypeError, match="incompatible function arguments"):
+        m.round_trip_view_tensor(tensor_ref.tolist())
diff --git a/tests/test_numpy_array.cpp b/tests/test_numpy_array.cpp
@@ -586,4 +586,13 @@ TEST_SUBMODULE(numpy_array, sm) {
     sm.def("return_array_pyobject_ptr_from_list", return_array_from_list<PyObject *>);
     sm.def("return_array_handle_from_list", return_array_from_list<py::handle>);
     sm.def("return_array_object_from_list", return_array_from_list<py::object>);
+
+    sm.def(
+        "round_trip_array_t",
+        [](const py::array_t<float> &x) -> py::array_t<float> { return x; },
+        py::arg("x"));
+    sm.def(
+        "round_trip_array_t_noconvert",
+        [](const py::array_t<float> &x) -> py::array_t<float> { return x; },
+        py::arg("x").noconvert());
 }
diff --git a/tests/test_numpy_array.py b/tests/test_numpy_array.py
@@ -687,3 +687,17 @@ def test_return_array_object_cpp_loop(return_array, unwrap):
     assert isinstance(arr_from_list, np.ndarray)
     assert arr_from_list.dtype == np.dtype("O")
     assert unwrap(arr_from_list) == [6, "seven", -8.0]
+
+
+def test_arraylike_signature(doc):
+    assert (
+        doc(m.round_trip_array_t)
+        == "round_trip_array_t(x: typing.Annotated[numpy.typing.ArrayLike, numpy.float32]) -> numpy.typing.NDArray[numpy.float32]"
+    )
+    assert (
+        doc(m.round_trip_array_t_noconvert)
+        == "round_trip_array_t_noconvert(x: numpy.typing.NDArray[numpy.float32]) -> numpy.typing.NDArray[numpy.float32]"
+    )
+    m.round_trip_array_t([1, 2, 3])
+    with pytest.raises(TypeError, match="incompatible function arguments"):
+        m.round_trip_array_t_noconvert([1, 2, 3])