eigen: Disable dtype=object arrays from being referenced

EricCousineau-TRI · EricCousineau-TRI · commit 719532b522e9 · 2018-05-23T10:55:26.000-04:00
diff --git a/include/pybind11/eigen.h b/include/pybind11/eigen.h
@@ -113,9 +113,8 @@ struct eigen_extract_stride<Eigen::Map<PlainObjectType, MapOptions, StrideType>>
 template <typename PlainObjectType, int Options, typename StrideType>
 struct eigen_extract_stride<Eigen::Ref<PlainObjectType, Options, StrideType>> { using type = StrideType; };
 
-template <typename Scalar> bool is_pyobject_() {
-    return static_cast<pybind11::detail::npy_api::constants>(npy_format_descriptor<Scalar>::value) == npy_api::NPY_OBJECT_;
-}
+template <typename Scalar>
+using is_pyobject_dtype = std::is_base_of<npy_format_descriptor_object, npy_format_descriptor<Scalar>>;
 
 // Helper struct for extracting information from an Eigen type
 template <typename Type_> struct EigenProps {
@@ -149,9 +148,7 @@ template <typename Type_> struct EigenProps {
         const auto dims = a.ndim();
         if (dims < 1 || dims > 2)
             return false;
-        bool is_pyobject = false;
-        if (is_pyobject_<Scalar>())
-            is_pyobject = true;
+        constexpr bool is_pyobject = is_pyobject_dtype<Scalar>::value;
         ssize_t scalar_size = (is_pyobject ? static_cast<ssize_t>(sizeof(PyObject*)) :
                                static_cast<ssize_t>(sizeof(Scalar)));
         if (dims == 2) { // Matrix type: require exact match (or dynamic)
@@ -233,17 +230,27 @@ template <typename props> handle eigen_array_cast(typename props::Type const &sr
                       src.data(), base);
     }
     else {
+        if (base) {
+            // Should be disabled by upstream calls to this method.
+            // TODO(eric.cousineau): Write tests to ensure that this is not
+            // reachable.
+            throw cast_error(
+                "dtype=object does not permit array referencing. "
+                "(NOTE: this generally not be reachable, as upstream APIs "
+                "should fail before this.");
+        }
+        handle empty_base{};
+        auto policy = return_value_policy::copy;
         if (props::vector) {
             a = array(
                 npy_format_descriptor<Scalar>::dtype(),
                 { (size_t) src.size() },
                 nullptr,
-                base
+                empty_base
             );
-            auto policy = base ? return_value_policy::automatic_reference : return_value_policy::copy;
             for (ssize_t i = 0; i < src.size(); ++i) {
                 const Scalar src_val = props::fixed_rows ? src(0, i) : src(i, 0);
-                auto value_ = reinterpret_steal<object>(make_caster<Scalar>::cast(src_val, policy, base));
+                auto value_ = reinterpret_steal<object>(make_caster<Scalar>::cast(src_val, policy, empty_base));
                 if (!value_)
                     return handle();
                 a.attr("itemset")(i, value_);
@@ -254,12 +261,11 @@ template <typename props> handle eigen_array_cast(typename props::Type const &sr
                 npy_format_descriptor<Scalar>::dtype(),
                 {(size_t) src.rows(), (size_t) src.cols()},
                 nullptr,
-                base
+                empty_base
             );
-            auto policy = base ? return_value_policy::automatic_reference : return_value_policy::copy;
             for (ssize_t i = 0; i < src.rows(); ++i) {
                 for (ssize_t j = 0; j < src.cols(); ++j) {
-                    auto value_ = reinterpret_steal<object>(make_caster<Scalar>::cast(src(i, j), policy, base));
+                    auto value_ = reinterpret_steal<object>(make_caster<Scalar>::cast(src(i, j), policy, empty_base));
                     if (!value_)
                         return handle();
                     a.attr("itemset")(i, j, value_);
@@ -323,7 +329,7 @@ struct type_caster<Type, enable_if_t<is_eigen_dense_plain<Type>::value>> {
         int result = 0;
         // Allocate the new type, then build a numpy reference into it
         value = Type(fits.rows, fits.cols);
-        bool is_pyobject = is_pyobject_<Scalar>();
+        constexpr bool is_pyobject = is_pyobject_dtype<Scalar>::value;
 
         if (!is_pyobject) {
             auto ref = reinterpret_steal<array>(eigen_ref_array<props>(value));
@@ -374,22 +380,40 @@ struct type_caster<Type, enable_if_t<is_eigen_dense_plain<Type>::value>> {
     // Cast implementation
     template <typename CType>
     static handle cast_impl(CType *src, return_value_policy policy, handle parent) {
-        switch (policy) {
-            case return_value_policy::take_ownership:
-            case return_value_policy::automatic:
-                return eigen_encapsulate<props>(src);
-            case return_value_policy::move:
-                return eigen_encapsulate<props>(new CType(std::move(*src)));
-            case return_value_policy::copy:
-                return eigen_array_cast<props>(*src);
-            case return_value_policy::reference:
-            case return_value_policy::automatic_reference:
-                return eigen_ref_array<props>(*src);
-            case return_value_policy::reference_internal:
-                return eigen_ref_array<props>(*src, parent);
-            default:
-                throw cast_error("unhandled return_value_policy: should not happen!");
-        };
+        constexpr bool is_pyobject = is_pyobject_dtype<Scalar>::value;
+        if (!is_pyobject) {
+            switch (policy) {
+                case return_value_policy::take_ownership:
+                case return_value_policy::automatic:
+                    return eigen_encapsulate<props>(src);
+                case return_value_policy::move:
+                    return eigen_encapsulate<props>(new CType(std::move(*src)));
+                case return_value_policy::copy:
+                    return eigen_array_cast<props>(*src);
+                case return_value_policy::reference:
+                case return_value_policy::automatic_reference:
+                    return eigen_ref_array<props>(*src);
+                case return_value_policy::reference_internal:
+                    return eigen_ref_array<props>(*src, parent);
+                default:
+                    throw cast_error("unhandled return_value_policy: should not happen!");
+            };
+        } else {
+            // For arrays of `dtype=object`, referencing is invalid, so we should squash that as soon as possible.
+            switch (policy) {
+                case return_value_policy::automatic:
+                case return_value_policy::move:
+                case return_value_policy::copy:
+                case return_value_policy::automatic_reference:
+                    return eigen_array_cast<props>(*src);
+                case return_value_policy::take_ownership:
+                case return_value_policy::reference:
+                case return_value_policy::reference_internal:
+                    throw cast_error("dtype=object arrays must be copied, and cannot be referenced");
+                default:
+                    throw cast_error("unhandled return_value_policy: should not happen!");
+            };
+        }
     }
 
 public:
@@ -446,6 +470,7 @@ struct return_value_policy_override<Return, enable_if_t<is_eigen_dense_map<Retur
 template <typename MapType> struct eigen_map_caster {
 private:
     using props = EigenProps<MapType>;
+    using Scalar = typename props::Scalar;
 
 public:
 
@@ -456,18 +481,33 @@ template <typename MapType> struct eigen_map_caster {
     // that this means you need to ensure you don't destroy the object in some other way (e.g. with
     // an appropriate keep_alive, or with a reference to a statically allocated matrix).
     static handle cast(const MapType &src, return_value_policy policy, handle parent) {
-        switch (policy) {
-            case return_value_policy::copy:
-                return eigen_array_cast<props>(src);
-            case return_value_policy::reference_internal:
-                return eigen_array_cast<props>(src, parent, is_eigen_mutable_map<MapType>::value);
-            case return_value_policy::reference:
-            case return_value_policy::automatic:
-            case return_value_policy::automatic_reference:
-                return eigen_array_cast<props>(src, none(), is_eigen_mutable_map<MapType>::value);
-            default:
-                // move, take_ownership don't make any sense for a ref/map:
-                pybind11_fail("Invalid return_value_policy for Eigen Map/Ref/Block type");
+        if (!is_pyobject_dtype<Scalar>::value) {
+            switch (policy) {
+                case return_value_policy::copy:
+                    return eigen_array_cast<props>(src);
+                case return_value_policy::reference_internal:
+                    return eigen_array_cast<props>(src, parent, is_eigen_mutable_map<MapType>::value);
+                case return_value_policy::reference:
+                case return_value_policy::automatic:
+                case return_value_policy::automatic_reference:
+                    return eigen_array_cast<props>(src, none(), is_eigen_mutable_map<MapType>::value);
+                default:
+                    // move, take_ownership don't make any sense for a ref/map:
+                    pybind11_fail("Invalid return_value_policy for Eigen Map/Ref/Block type");
+            }
+        } else {
+            switch (policy) {
+                case return_value_policy::copy:
+                    return eigen_array_cast<props>(src);
+                case return_value_policy::reference_internal:
+                case return_value_policy::reference:
+                case return_value_policy::automatic:
+                case return_value_policy::automatic_reference:
+                    throw cast_error("dtype=object arrays must be copied, and cannot be referenced");
+                default:
+                    // move, take_ownership don't make any sense for a ref/map:
+                    pybind11_fail("Invalid return_value_policy for Eigen Map/Ref/Block type");
+            }
         }
     }
 
@@ -519,9 +559,14 @@ struct type_caster<
         bool need_copy = !isinstance<Array>(src);
 
         EigenConformable<props::row_major> fits;
-        bool is_pyobject = false;
-        if (is_pyobject_<Scalar>()) {
-            is_pyobject = true;
+        constexpr bool is_pyobject = is_pyobject_dtype<Scalar>::value;
+        // TODO(eric.cousineau): Make this compile-time once Drake does not use this in any code
+        // for scalar types.
+        //static_assert(!(is_pyobject && need_writeable), "dtype=object cannot provide writeable references");
+        if (is_pyobject && need_writeable) {
+            throw cast_error("dtype=object cannot provide writeable references");
+        }
+        if (is_pyobject) {
             need_copy = true;
         }
         if (!need_copy) {
diff --git a/include/pybind11/numpy.h b/include/pybind11/numpy.h
@@ -1230,19 +1230,22 @@ template <typename T, typename SFINAE> struct npy_format_descriptor {
         (::std::vector<::pybind11::detail::field_descriptor> \
          {PYBIND11_MAP2_LIST (PYBIND11_FIELD_DESCRIPTOR_EX, Type, __VA_ARGS__)})
 
+struct npy_format_descriptor_object {
+public:
+    enum { value = npy_api::NPY_OBJECT_ };
+    static pybind11::dtype dtype() {
+        if (auto ptr = npy_api::get().PyArray_DescrFromType_(value)) {
+            return reinterpret_borrow<pybind11::dtype>(ptr);
+        }
+        pybind11_fail("Unsupported buffer format!");
+    }
+    static constexpr auto name = _("object");
+};
+
 #define PYBIND11_NUMPY_OBJECT_DTYPE(Type) \
     namespace pybind11 { namespace detail { \
-        template <> struct npy_format_descriptor<Type> { \
-        public: \
-            enum { value = npy_api::NPY_OBJECT_ }; \
-            static pybind11::dtype dtype() { \
-                if (auto ptr = npy_api::get().PyArray_DescrFromType_(value)) { \
-                    return reinterpret_borrow<pybind11::dtype>(ptr); \
-                } \
-                pybind11_fail("Unsupported buffer format!"); \
-            } \
-            static constexpr auto name = _("object"); \
-        }; \
+        template <> struct npy_format_descriptor<Type> : \
+            public npy_format_descriptor_object {}; \
     }}
 
 #endif // __CLION_IDE__
diff --git a/tests/test_eigen.cpp b/tests/test_eigen.cpp
@@ -61,6 +61,16 @@ void reset_refs() {
     reset_ref(get_rm());
 }
 
+VectorXADScalar& get_cm_adscalar() {
+    static VectorXADScalar value(1);
+    return value;
+};
+VectorXADScalarR& get_rm_adscalar() {
+    static VectorXADScalarR value(1);
+    return value;
+};
+
+
 // Returns element 2,1 from a matrix (used to test copy/nocopy)
 double get_elem(Eigen::Ref<const Eigen::MatrixXd> m) { return m(2, 1); };
 
@@ -106,9 +116,7 @@ TEST_SUBMODULE(eigen, m) {
     m.def("double_adscalar_row", [](const VectorXADScalarR &x) -> VectorXADScalarR { return 2.0f * x; });
     m.def("double_complex", [](const Eigen::VectorXcf &x) -> Eigen::VectorXcf { return 2.0f * x; });
     m.def("double_threec", [](py::EigenDRef<Eigen::Vector3f> x) { x *= 2; });
-    m.def("double_adscalarc", [](py::EigenDRef<VectorXADScalar> x) { x *= 2; });
     m.def("double_threer", [](py::EigenDRef<Eigen::RowVector3f> x) { x *= 2; });
-    m.def("double_adscalarr", [](py::EigenDRef<VectorXADScalarR> x) { x *= 2; });
     m.def("double_mat_cm", [](Eigen::MatrixXf x) -> Eigen::MatrixXf { return 2.0f * x; });
     m.def("double_mat_rm", [](DenseMatrixR x) -> DenseMatrixR { return 2.0f * x; });
 
@@ -130,6 +138,8 @@ TEST_SUBMODULE(eigen, m) {
     // Mutators (Eigen maps into numpy variables):
     m.def("add_rm", add_rm); // Only takes row-contiguous
     m.def("add_cm", add_cm); // Only takes column-contiguous
+    m.def("add_rm_adscalar", [](py::EigenDRef<VectorXADScalarR> x) { x.array() += 2; });
+    m.def("add_cm_adscalar", [](py::EigenDRef<VectorXADScalar> x) { x.array() += 2; });
     // Overloaded versions that will accept either row or column contiguous:
     m.def("add1", add_rm);
     m.def("add1", add_cm);
@@ -141,9 +151,17 @@ TEST_SUBMODULE(eigen, m) {
     // Return mutable references (numpy maps into eigen variables)
     m.def("get_cm_ref", []() { return Eigen::Ref<Eigen::MatrixXd>(get_cm()); });
     m.def("get_rm_ref", []() { return Eigen::Ref<MatrixXdR>(get_rm()); });
+    m.def("get_cm_ref_adscalar", []() {
+        return py::EigenDRef<VectorXADScalar>(get_cm_adscalar());
+    });
+    m.def("get_rm_ref_adscalar", []() {
+        return py::EigenDRef<VectorXADScalarR>(get_rm_adscalar());
+    });
     // The same references, but non-mutable (numpy maps into eigen variables, but is !writeable)
     m.def("get_cm_const_ref", []() { return Eigen::Ref<const Eigen::MatrixXd>(get_cm()); });
     m.def("get_rm_const_ref", []() { return Eigen::Ref<const MatrixXdR>(get_rm()); });
+    m.def("get_cm_const_ref_adscalar", []() { return Eigen::Ref<const VectorXADScalar>(get_cm_adscalar()); });
+    m.def("get_rm_const_ref_adscalar", []() { return Eigen::Ref<const VectorXADScalarR>(get_rm_adscalar()); });
 
     m.def("reset_refs", reset_refs); // Restores get_{cm,rm}_ref to original values
 
@@ -153,11 +171,12 @@ TEST_SUBMODULE(eigen, m) {
         return m;
     }, py::return_value_policy::reference);
 
-    // Increments ADScalar Matrix
-    m.def("incr_adscalar_matrix", [](Eigen::Ref<DenseADScalarMatrixC> m, double v) {
-      m += DenseADScalarMatrixC::Constant(m.rows(), m.cols(), v);
-      return m;
-    }, py::return_value_policy::reference);
+    // Increments ADScalar Matrix, returns a copy.
+    m.def("incr_adscalar_matrix", [](const Eigen::Ref<const DenseADScalarMatrixC>& m, double v) {
+      DenseADScalarMatrixC out = m;
+      out.array() += v;
+      return out;
+    });
 
     // Same, but accepts a matrix of any strides
     m.def("incr_matrix_any", [](py::EigenDRef<Eigen::MatrixXd> m, double v) {
@@ -347,10 +366,7 @@ TEST_SUBMODULE(eigen, m) {
     m.def("cpp_matrix_shape", [](const MatrixX<ADScalar>& A) {
         return py::make_tuple(A.rows(), A.cols());
     });
-    // TODO(eric.cousineau): Unless `dtype=ADScalar` (user-defined) and not
-    // `dtype=object`, we should kill any usages of `Eigen::Ref<>` or any
-    // const-references.
-    m.def("cpp_matrix_shape_ref", [](const Eigen::Ref<MatrixX<ADScalar>>& A) {
+    m.def("cpp_matrix_shape_ref", [](const Eigen::Ref<const MatrixX<ADScalar>>& A) {
         return py::make_tuple(A.rows(), A.cols());
     });
 
diff --git a/tests/test_eigen.py b/tests/test_eigen.py
@@ -182,15 +182,31 @@ def test_eigen_passing_adscalar():
     incremented_adscalar_mat = conv_double_to_adscalar(m.incr_adscalar_matrix(adscalar_mat, 7.),
                                                        vice_versa=True)
     np.testing.assert_array_equal(incremented_adscalar_mat, ref + 7)
-    # The original adscalar_mat remains unchanged in spite of passing by reference.
+    # The original adscalar_mat remains unchanged in spite of passing by reference, since
+    # `Eigen::Ref<const CType>` permits copying, and copying is the only valid operation for
+    # `dtype=object`.
     np.testing.assert_array_equal(conv_double_to_adscalar(adscalar_mat, vice_versa=True), ref)
 
-    # Changes in Python are not reflected in C++ when internal_reference is returned
+    # Changes in Python are not reflected in C++ when internal_reference is returned.
+    # These conversions should be disabled at runtime.
+
+    def expect_ref_error(func):
+        with pytest.raises(RuntimeError) as excinfo:
+            func()
+        assert "dtype=object" in str(excinfo)
+        assert "reachable" not in str(excinfo)
+
+    # - Return arguments.
+    expect_ref_error(lambda: m.get_cm_ref_adscalar())
+    expect_ref_error(lambda: m.get_rm_ref_adscalar())
+    expect_ref_error(lambda: m.get_cm_const_ref_adscalar())
+    expect_ref_error(lambda: m.get_rm_const_ref_adscalar())
+    # - - Mutable lvalues referenced via `reference_internal`.
     return_tester = m.ReturnTester()
-    a = return_tester.get_ADScalarMat()
-    a[1, 1] = m.AutoDiffXd(4, np.ones(1))
-    b = return_tester.get_ADScalarMat()
-    assert(np.isclose(b[1, 1].value(), 7.))
+    expect_ref_error(lambda: return_tester.get_ADScalarMat())
+    # - Input arguments, writeable `Ref<>`s.
+    expect_ref_error(lambda: m.add_cm_adscalar(adscalar_vec_col))
+    expect_ref_error(lambda: m.add_rm_adscalar(adscalar_vec_row))
 
     # Checking Issue 1105
     assert m.iss1105_col_obj(adscalar_vec_col[:, None])
