Adding tests & minor bugfix

tdegeus · tdegeus · commit 428ed03b5dbc · 2021-06-05T16:17:21.000+02:00
diff --git a/include/xtensor-python/xtensor_type_caster_base.hpp b/include/xtensor-python/xtensor_type_caster_base.hpp
@@ -37,7 +37,7 @@ namespace pybind11
         {
             static auto get(handle src)
             {
-                return array_t<T, array::f_style>::ensure(src);
+                return array_t<T, array::f_style | array::forcecast>::ensure(src);
             }
         };
 
@@ -51,8 +51,7 @@ namespace pybind11
         {
             static auto get(handle src)
             {
-                auto buf = xtensor_get_buffer<T, L>::get(src);
-                return buf;
+                return xtensor_get_buffer<T, L>::get(src);
             }
         };
 
@@ -61,11 +60,7 @@ namespace pybind11
         {
             static auto get(handle src)
             {
-                auto buf = xtensor_get_buffer<T, L>::get(src);
-                if (buf.ndim() != N) {
-                    return false;
-                }
-                return buf;
+                return xtensor_get_buffer<T, L>::get(src);
             }
         };
 
@@ -98,6 +93,27 @@ namespace pybind11
         };
 
 
+        template <class T>
+        struct xtensor_verify
+        {
+            template <class B>
+            static bool get(const B& buf)
+            {
+                return true;
+            }
+        };
+
+        template <class T, std::size_t N, xt::layout_type L>
+        struct xtensor_verify<xt::xtensor<T, N, L>>
+        {
+            template <class B>
+            static bool get(const B& buf)
+            {
+                return buf.ndim() == N;
+            }
+        };
+
+
         // Casts a strided expression type to numpy array.If given a base,
         // the numpy array references the src data, otherwise it'll make a copy.
         // The writeable attributes lets you specify writeable flag for the array.
@@ -192,11 +208,14 @@ namespace pybind11
                 if (!buf) {
                     return false;
                 }
+                if (!xtensor_verify<Type>::get(buf)) {
+                    return false;
+                }
 
                 std::vector<size_t> shape(buf.ndim());
                 std::copy(buf.shape(), buf.shape() + buf.ndim(), shape.begin());
-                value = Type(shape);
-                std::copy(buf.data(), buf.data() + buf.size(), value.begin());
+                value = Type::from_shape(shape);
+                std::copy(buf.data(), buf.data() + buf.size(), value.data());
 
                 return true;
             }
diff --git a/test_python/main.cpp b/test_python/main.cpp
@@ -33,6 +33,33 @@ xt::pyarray<double> example2(xt::pyarray<double>& m)
     return m + 2;
 }
 
+xt::xarray<int> example3_xarray(const xt::xarray<int>& m)
+{
+    return xt::transpose(m) + 2;
+}
+
+xt::xarray<int, xt::layout_type::column_major> example3_xarray_colmajor(
+    const xt::xarray<int, xt::layout_type::column_major>& m)
+{
+    return xt::transpose(m) + 2;
+}
+
+xt::xtensor<int, 3> example3_xtensor3(const xt::xtensor<int, 3>& m)
+{
+    return xt::transpose(m) + 2;
+}
+
+xt::xtensor<int, 2> example3_xtensor2(const xt::xtensor<int, 2>& m)
+{
+    return xt::transpose(m) + 2;
+}
+
+xt::xtensor<int, 2, xt::layout_type::column_major> example3_xtensor2_colmajor(
+    const xt::xtensor<int, 2, xt::layout_type::column_major>& m)
+{
+    return xt::transpose(m) + 2;
+}
+
 // Readme Examples
 
 double readme_example1(xt::pyarray<double>& m)
@@ -249,6 +276,11 @@ PYBIND11_MODULE(xtensor_python_test, m)
 
     m.def("example1", example1);
     m.def("example2", example2);
+    m.def("example3_xarray", example3_xarray);
+    m.def("example3_xarray_colmajor", example3_xarray_colmajor);
+    m.def("example3_xtensor3", example3_xtensor3);
+    m.def("example3_xtensor2", example3_xtensor2);
+    m.def("example3_xtensor2_colmajor", example3_xtensor2_colmajor);
 
     m.def("complex_overload", no_complex_overload);
     m.def("complex_overload", complex_overload);
diff --git a/test_python/test_pyarray.py b/test_python/test_pyarray.py
@@ -36,6 +36,22 @@ def test_example2(self):
         y = xt.example2(x)
         np.testing.assert_allclose(y, res, 1e-12)
 
+    def test_example3(self):
+        x = np.arange(2 * 3).reshape(2, 3)
+        xc = np.asfortranarray(x)
+        y = np.arange(2 * 3 * 4).reshape(2, 3, 4)
+        v = y[1:, 1:, 0]
+        z = np.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5)
+        np.testing.assert_array_equal(xt.example3_xarray(x), x.T + 2)
+        np.testing.assert_array_equal(xt.example3_xarray_colmajor(xc), xc.T + 2)
+        np.testing.assert_array_equal(xt.example3_xtensor3(y), y.T + 2)
+        np.testing.assert_array_equal(xt.example3_xtensor2(x), x.T + 2)
+        np.testing.assert_array_equal(xt.example3_xtensor2(y[1:, 1:, 0]), v.T + 2)
+        np.testing.assert_array_equal(xt.example3_xtensor2_colmajor(xc), xc.T + 2)
+
+        with self.assertRaises(TypeError):
+            xt.example3_xtensor3(x)
+
     def test_vectorize(self):
         x1 = np.array([[0, 1], [2, 3]])
         x2 = np.array([0, 1])
