Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft: np.typing+typing.Annotated using io_name #5502

Closed
wants to merge 2 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 14 additions & 9 deletions include/pybind11/eigen/matrix.h
Original file line number Diff line number Diff line change
Expand Up @@ -225,19 +225,22 @@ struct EigenProps {
= !show_c_contiguous && show_order && requires_col_major;

static constexpr auto descriptor
= const_name("numpy.ndarray[") + npy_format_descriptor<Scalar>::name + const_name("[")
= const_name("typing.Annotated[")
+ io_name("numpy.typing.ArrayLike, ", "numpy.typing.NDArray[")
+ npy_format_descriptor<Scalar>::name + io_name("", "]") + const_name(", \"[")
+ const_name<fixed_rows>(const_name<(size_t) rows>(), const_name("m")) + const_name(", ")
+ const_name<fixed_cols>(const_name<(size_t) cols>(), const_name("n")) + const_name("]")
+
+ const_name<fixed_cols>(const_name<(size_t) cols>(), const_name("n"))
+ const_name("]\"") +
// For a reference type (e.g. Ref<MatrixXd>) we have other constraints that might need to
// be satisfied: writeable=True (for a mutable reference), and, depending on the map's
// stride options, possibly f_contiguous or c_contiguous. We include them in the
// descriptor output to provide some hint as to why a TypeError is occurring (otherwise
// it can be confusing to see that a function accepts a 'numpy.ndarray[float64[3,2]]' and
// an error message that you *gave* a numpy.ndarray of the right type and dimensions.
const_name<show_writeable>(", flags.writeable", "")
+ const_name<show_c_contiguous>(", flags.c_contiguous", "")
+ const_name<show_f_contiguous>(", flags.f_contiguous", "") + const_name("]");
// it can be confusing to see that a function accepts a
// 'typing.Annotated[numpy.typing.NDArray[numpy.float64], "[3,2]"]' and an error message
// that you *gave* a numpy.ndarray of the right type and dimensions.
const_name<show_writeable>(", \"flags.writeable\"", "")
+ const_name<show_c_contiguous>(", \"flags.c_contiguous\"", "")
+ const_name<show_f_contiguous>(", \"flags.f_contiguous\"", "") + const_name("]");
};

// Casts an Eigen type to numpy array. If given a base, the numpy array references the src data,
Expand Down Expand Up @@ -441,7 +444,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
Expand Down
18 changes: 12 additions & 6 deletions include/pybind11/eigen/tensor.h
Original file line number Diff line number Diff line change
Expand Up @@ -124,13 +124,16 @@ struct eigen_tensor_helper<
template <typename Type, bool ShowDetails, bool NeedsWriteable = false>
struct get_tensor_descriptor {
static constexpr auto details
= const_name<NeedsWriteable>(", flags.writeable", "") + const_name
= const_name<NeedsWriteable>(", \"flags.writeable\"", "") + const_name
< static_cast<int>(Type::Layout)
== static_cast<int>(Eigen::RowMajor) > (", flags.c_contiguous", ", flags.f_contiguous");
== static_cast<int>(Eigen::RowMajor)
> (", \"flags.c_contiguous\"", ", \"flags.f_contiguous\"");
static constexpr auto value
= const_name("numpy.ndarray[") + npy_format_descriptor<typename Type::Scalar>::name
+ const_name("[") + eigen_tensor_helper<remove_cv_t<Type>>::dimensions_descriptor
+ const_name("]") + const_name<ShowDetails>(details, const_name("")) + const_name("]");
= const_name("typing.Annotated[")
+ io_name("numpy.typing.ArrayLike, ", "numpy.typing.NDArray[")
+ npy_format_descriptor<typename Type::Scalar>::name + io_name("", "]")
+ const_name(", \"[") + eigen_tensor_helper<remove_cv_t<Type>>::dimensions_descriptor
+ const_name("]\"") + const_name<ShowDetails>(details, const_name("")) + const_name("]");
};

// When EIGEN_AVOID_STL_ARRAY is defined, Eigen::DSizes<T, 0> does not have the begin() member
Expand Down Expand Up @@ -502,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); }
Expand Down
3 changes: 2 additions & 1 deletion include/pybind11/numpy.h
Original file line number Diff line number Diff line change
Expand Up @@ -2183,7 +2183,8 @@ vectorize_helper<Func, Return, Args...> vectorize_extractor(const Func &f, Retur
template <typename T, int Flags>
struct handle_type_name<array_t<T, Flags>> {
static constexpr auto name
= const_name("numpy.ndarray[") + npy_format_descriptor<T>::name + const_name("]");
= io_name("typing.Annotated[numpy.typing.ArrayLike, ", "numpy.typing.NDArray[")
+ npy_format_descriptor<T>::name + const_name("]");
};

PYBIND11_NAMESPACE_END(detail)
Expand Down
4 changes: 4 additions & 0 deletions tests/test_eigen_matrix.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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; });
}
60 changes: 40 additions & 20 deletions tests/test_eigen_matrix.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,19 +95,20 @@ def test_mutator_descriptors():
with pytest.raises(TypeError) as excinfo:
m.fixed_mutator_r(zc)
assert (
"(arg0: numpy.ndarray[numpy.float32[5, 6],"
" flags.writeable, flags.c_contiguous]) -> None" in str(excinfo.value)
'(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: numpy.ndarray[numpy.float32[5, 6],"
" flags.writeable, flags.f_contiguous]) -> None" in str(excinfo.value)
'(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: numpy.ndarray[numpy.float32[5, 6], flags.writeable]) -> None" in str(
excinfo.value
assert (
'(arg0: typing.Annotated[numpy.typing.NDArray[numpy.float32], "[5, 6]", "flags.writeable"]) -> None'
in str(excinfo.value)
)
zr.flags.writeable = False
with pytest.raises(TypeError):
Expand Down Expand Up @@ -201,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: numpy.ndarray[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"))
Expand All @@ -213,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: numpy.ndarray[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"))
Expand Down Expand Up @@ -634,37 +635,37 @@ def test_nocopy_wrapper():
with pytest.raises(TypeError) as excinfo:
m.get_elem_nocopy(int_matrix_colmajor)
assert "get_elem_nocopy(): incompatible function arguments." in str(excinfo.value)
assert ", flags.f_contiguous" in str(excinfo.value)
assert ', "flags.f_contiguous"' in str(excinfo.value)
assert m.get_elem_nocopy(dbl_matrix_colmajor) == 8
with pytest.raises(TypeError) as excinfo:
m.get_elem_nocopy(int_matrix_rowmajor)
assert "get_elem_nocopy(): incompatible function arguments." in str(excinfo.value)
assert ", flags.f_contiguous" in str(excinfo.value)
assert ', "flags.f_contiguous"' in str(excinfo.value)
with pytest.raises(TypeError) as excinfo:
m.get_elem_nocopy(dbl_matrix_rowmajor)
assert "get_elem_nocopy(): incompatible function arguments." in str(excinfo.value)
assert ", flags.f_contiguous" in str(excinfo.value)
assert ', "flags.f_contiguous"' in str(excinfo.value)

# For the row-major test, we take a long matrix in row-major, so only the third is allowed:
with pytest.raises(TypeError) as excinfo:
m.get_elem_rm_nocopy(int_matrix_colmajor)
assert "get_elem_rm_nocopy(): incompatible function arguments." in str(
excinfo.value
)
assert ", flags.c_contiguous" in str(excinfo.value)
assert ', "flags.c_contiguous"' in str(excinfo.value)
with pytest.raises(TypeError) as excinfo:
m.get_elem_rm_nocopy(dbl_matrix_colmajor)
assert "get_elem_rm_nocopy(): incompatible function arguments." in str(
excinfo.value
)
assert ", flags.c_contiguous" in str(excinfo.value)
assert ', "flags.c_contiguous"' in str(excinfo.value)
assert m.get_elem_rm_nocopy(int_matrix_rowmajor) == 8
with pytest.raises(TypeError) as excinfo:
m.get_elem_rm_nocopy(dbl_matrix_rowmajor)
assert "get_elem_rm_nocopy(): incompatible function arguments." in str(
excinfo.value
)
assert ", flags.c_contiguous" in str(excinfo.value)
assert ', "flags.c_contiguous"' in str(excinfo.value)


def test_eigen_ref_life_support():
Expand Down Expand Up @@ -700,25 +701,25 @@ def test_dense_signature(doc):
assert (
doc(m.double_col)
== """
double_col(arg0: numpy.ndarray[numpy.float32[m, 1]]) -> numpy.ndarray[numpy.float32[m, 1]]
double_col(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[m, 1]"]) -> typing.Annotated[numpy.typing.NDArray[numpy.float32], "[m, 1]"]
"""
)
assert (
doc(m.double_row)
== """
double_row(arg0: numpy.ndarray[numpy.float32[1, n]]) -> numpy.ndarray[numpy.float32[1, n]]
double_row(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[1, n]"]) -> typing.Annotated[numpy.typing.NDArray[numpy.float32], "[1, n]"]
"""
)
assert doc(m.double_complex) == (
"""
double_complex(arg0: numpy.ndarray[numpy.complex64[m, 1]])"""
""" -> numpy.ndarray[numpy.complex64[m, 1]]
double_complex(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.complex64, "[m, 1]"])"""
""" -> typing.Annotated[numpy.typing.NDArray[numpy.complex64], "[m, 1]"]
"""
)
assert doc(m.double_mat_rm) == (
"""
double_mat_rm(arg0: numpy.ndarray[numpy.float32[m, n]])"""
""" -> numpy.ndarray[numpy.float32[m, n]]
double_mat_rm(arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32, "[m, n]"])"""
""" -> typing.Annotated[numpy.typing.NDArray[numpy.float32], "[m, n]"]
"""
)

Expand Down Expand Up @@ -817,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]])
39 changes: 31 additions & 8 deletions tests/test_eigen_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -271,23 +271,46 @@ def test_round_trip_references_actually_refer(m):
@pytest.mark.parametrize("m", submodules)
def test_doc_string(m, doc):
assert (
doc(m.copy_tensor) == "copy_tensor() -> numpy.ndarray[numpy.float64[?, ?, ?]]"
doc(m.copy_tensor)
== 'copy_tensor() -> typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]"]'
)
assert (
doc(m.copy_fixed_tensor)
== "copy_fixed_tensor() -> numpy.ndarray[numpy.float64[3, 5, 2]]"
== 'copy_fixed_tensor() -> typing.Annotated[numpy.typing.NDArray[numpy.float64], "[3, 5, 2]"]'
)
assert (
doc(m.reference_const_tensor)
== "reference_const_tensor() -> numpy.ndarray[numpy.float64[?, ?, ?]]"
== 'reference_const_tensor() -> typing.Annotated[numpy.typing.NDArray[numpy.float64], "[?, ?, ?]"]'
)

order_flag = f"flags.{m.needed_options.lower()}_contiguous"
order_flag = f'"flags.{m.needed_options.lower()}_contiguous"'
assert doc(m.round_trip_view_tensor) == (
f"round_trip_view_tensor(arg0: numpy.ndarray[numpy.float64[?, ?, ?], flags.writeable, {order_flag}])"
f" -> numpy.ndarray[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: numpy.ndarray[numpy.float64[?, ?, ?], {order_flag}])"
" -> numpy.ndarray[numpy.float64[?, ?, ?]]"
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())
9 changes: 9 additions & 0 deletions tests/test_numpy_array.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
34 changes: 24 additions & 10 deletions tests/test_numpy_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -321,13 +321,13 @@ def test_overload_resolution(msg):
msg(excinfo.value)
== """
overloaded(): incompatible function arguments. The following argument types are supported:
1. (arg0: numpy.ndarray[numpy.float64]) -> str
2. (arg0: numpy.ndarray[numpy.float32]) -> str
3. (arg0: numpy.ndarray[numpy.int32]) -> str
4. (arg0: numpy.ndarray[numpy.uint16]) -> str
5. (arg0: numpy.ndarray[numpy.int64]) -> str
6. (arg0: numpy.ndarray[numpy.complex128]) -> str
7. (arg0: numpy.ndarray[numpy.complex64]) -> str
1. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float64]) -> str
2. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.float32]) -> str
3. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.int32]) -> str
4. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.uint16]) -> str
5. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.int64]) -> str
6. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.complex128]) -> str
7. (arg0: typing.Annotated[numpy.typing.ArrayLike, numpy.complex64]) -> str

Invoked with: 'not an array'
"""
Expand All @@ -343,8 +343,8 @@ def test_overload_resolution(msg):
assert m.overloaded3(np.array([1], dtype="intc")) == "int"
expected_exc = """
overloaded3(): incompatible function arguments. The following argument types are supported:
1. (arg0: numpy.ndarray[numpy.int32]) -> str
2. (arg0: numpy.ndarray[numpy.float64]) -> str
1. (arg0: numpy.typing.NDArray[numpy.int32]) -> str
2. (arg0: numpy.typing.NDArray[numpy.float64]) -> str

Invoked with: """

Expand Down Expand Up @@ -528,7 +528,7 @@ def test_index_using_ellipsis():
],
)
def test_format_descriptors_for_floating_point_types(test_func):
assert "numpy.ndarray[numpy.float" in test_func.__doc__
assert "typing.Annotated[numpy.typing.ArrayLike, numpy.float" in test_func.__doc__


@pytest.mark.parametrize("forcecast", [False, True])
Expand Down Expand Up @@ -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])
2 changes: 1 addition & 1 deletion tests/test_numpy_dtypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -373,7 +373,7 @@ def test_complex_array():
def test_signature(doc):
assert (
doc(m.create_rec_nested)
== "create_rec_nested(arg0: int) -> numpy.ndarray[NestedStruct]"
== "create_rec_nested(arg0: int) -> numpy.typing.NDArray[NestedStruct]"
)


Expand Down
Loading
Loading