Ensure the Eigen type_casters do not segfault when loading arrays with dtype=object

This commit is contained in:
Ralf W. Grosse-Kunstleve 2023-05-18 01:37:04 -07:00
parent 82ce80fae1
commit 20b9baf270
6 changed files with 69 additions and 0 deletions

View File

@ -290,6 +290,11 @@ struct type_caster<Type, enable_if_t<is_eigen_dense_plain<Type>::value>> {
using props = EigenProps<Type>;
bool load(handle src, bool convert) {
// dtype=object is not supported. See #1152 & #2259 for related experiments.
if (is_same_ignoring_cvref<Scalar, PyObject *>::value) {
return false;
}
// If we're in no-convert mode, only load if given an array of the correct type
if (!convert && !isinstance<array_t<Scalar>>(src)) {
return false;
@ -480,6 +485,11 @@ private:
public:
bool load(handle src, bool convert) {
// dtype=object is not supported. See #1152 & #2259 for related experiments.
if (is_same_ignoring_cvref<Scalar, PyObject *>::value) {
return false;
}
// First check whether what we have is already an array of the right type. If not, we
// can't avoid a copy (because the copy is also going to do type conversion).
bool need_copy = !isinstance<Array>(src);
@ -637,6 +647,11 @@ struct type_caster<Type, enable_if_t<is_eigen_sparse<Type>::value>> {
static constexpr bool rowMajor = Type::IsRowMajor;
bool load(handle src, bool) {
// dtype=object is not supported. See #1152 & #2259 for related experiments.
if (is_same_ignoring_cvref<Scalar, PyObject *>::value) {
return false;
}
if (!src) {
return false;
}

View File

@ -169,6 +169,11 @@ struct type_caster<Type, typename eigen_tensor_helper<Type>::ValidType> {
PYBIND11_TYPE_CASTER(Type, temp_name);
bool load(handle src, bool convert) {
// dtype=object is not supported. See #1152 & #2259 for related experiments.
if (is_same_ignoring_cvref<typename Type::Scalar, PyObject *>::value) {
return false;
}
if (!convert) {
if (!isinstance<array>(src)) {
return false;
@ -363,6 +368,11 @@ struct type_caster<Eigen::TensorMap<Type, Options>,
using Helper = eigen_tensor_helper<remove_cv_t<Type>>;
bool load(handle src, bool /*convert*/) {
// dtype=object is not supported. See #1152 & #2259 for related experiments.
if (is_same_ignoring_cvref<typename Type::Scalar, PyObject *>::value) {
return false;
}
// Note that we have a lot more checks here as we want to make sure to avoid copies
if (!isinstance<array>(src)) {
return false;

View File

@ -425,4 +425,10 @@ TEST_SUBMODULE(eigen_matrix, m) {
py::module_::import("numpy").attr("ones")(10);
return v[0](5);
});
m.def("pass_eigen_matrix_dtype_object",
[](const Eigen::Matrix<PyObject *, Eigen::Dynamic, Eigen::Dynamic> &) {});
m.def("pass_eigen_ref_matrix_dtype_object",
[](const Eigen::Ref<Eigen::Matrix<PyObject *, Eigen::Dynamic, Eigen::Dynamic>> &) {});
m.def("pass_eigen_sparse_matrix_dtype_object", [](const Eigen::SparseMatrix<PyObject *> &) {});
}

View File

@ -805,3 +805,20 @@ 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)
@pytest.mark.parametrize(
"pass_eigen_type_dtype_object",
[
m.pass_eigen_matrix_dtype_object,
m.pass_eigen_ref_matrix_dtype_object,
m.pass_eigen_sparse_matrix_dtype_object,
],
)
def test_pass_array_with_dtype_object(pass_eigen_type_dtype_object):
# Only the dtype matters (not shape etc.): dtype=object is (should be) the
# first check in the type_caster load() implementations.
obj = np.array([], dtype=object)
with pytest.raises(TypeError) as excinfo:
pass_eigen_type_dtype_object(obj)
assert "incompatible function arguments" in str(excinfo.value)

View File

@ -320,6 +320,10 @@ void init_tensor_module(pybind11::module &m) {
"round_trip_rank_0_view",
[](Eigen::TensorMap<Eigen::Tensor<double, 0, Options>> &tensor) { return tensor; },
py::return_value_policy::reference);
m.def("pass_eigen_tensor_dtype_object", [](const Eigen::Tensor<PyObject *, 0, Options> &) {});
m.def("pass_eigen_tensor_map_dtype_object",
[](Eigen::TensorMap<Eigen::Tensor<PyObject *, 0, Options>> &) {});
}
void test_module(py::module_ &m) {

View File

@ -286,3 +286,20 @@ def test_doc_string(m, doc):
f"round_trip_const_view_tensor(arg0: numpy.ndarray[numpy.float64[?, ?, ?], {order_flag}])"
" -> numpy.ndarray[numpy.float64[?, ?, ?]]"
)
@pytest.mark.parametrize("m", submodules)
@pytest.mark.parametrize(
"func_name",
[
"pass_eigen_tensor_dtype_object",
"pass_eigen_tensor_map_dtype_object",
],
)
def test_pass_array_with_dtype_object(m, func_name):
# Only the dtype matters (not shape etc.): dtype=object is (should be) the
# first check in the type_caster load() implementations.
obj = np.array([], dtype=object)
with pytest.raises(TypeError) as excinfo:
getattr(m, func_name)(obj)
assert "incompatible function arguments" in str(excinfo.value)