add test for complex overload (#87)

* add test for complex overload

* add non failing complex overload test

* fix complex overload loading

Author: wolfv

include/xtensor-python/pyarray.hpp

@@ -42,8 +42,20 @@ namespace pybind11
         {
             using type = xt::pyarray<T>;
 
-            bool load(handle src, bool)
+            bool load(handle src, bool convert)
             {
+                if (!convert)
+                {
+                    if (!PyArray_Check(src.ptr()))
+                    {
+                        return false;
+                    }
+                    int type_num = xt::detail::numpy_traits<T>::type_num;
+                    if (PyArray_TYPE(reinterpret_cast<PyArrayObject*>(src.ptr())) != type_num)
+                    {
+                        return false;
+                    }
+                }
                 value = type::ensure(src);
                 return static_cast<bool>(value);
             }

include/xtensor-python/pytensor.hpp

@@ -43,8 +43,21 @@ namespace pybind11
         {
             using type = xt::pytensor<T, N>;
 
-            bool load(handle src, bool)
+            bool load(handle src, bool convert)
             {
+                if (!convert)
+                {
+                    if (!PyArray_Check(src.ptr()))
+                    {
+                        return false;
+                    }
+                    int type_num = xt::detail::numpy_traits<T>::type_num;
+                    if (PyArray_TYPE(reinterpret_cast<PyArrayObject*>(src.ptr())) != type_num)
+                    {
+                        return false;
+                    }
+                }
+
                 value = type::ensure(src);
                 return static_cast<bool>(value);
             }

test_python/main.cpp

@@ -42,6 +42,25 @@ double readme_example2(double i, double j)
     return std::sin(i) -  std::cos(j);
 }
 
+auto complex_overload(const xt::pyarray<std::complex<double>>& a)
+{
+    return a;
+}
+auto no_complex_overload(const xt::pyarray<double>& a)
+{
+    return a;
+}
+
+auto complex_overload_reg(const std::complex<double>& a)
+{
+    return a;
+}
+
+auto no_complex_overload_reg(const double& a)
+{
+    return a;
+}
+
 // Vectorize Examples
 
 int add(int i, int j)
@@ -58,6 +77,11 @@ PYBIND11_PLUGIN(xtensor_python_test)
     m.def("example1", example1);
     m.def("example2", example2);
 
+    m.def("complex_overload", no_complex_overload);
+    m.def("complex_overload", complex_overload);
+    m.def("complex_overload_reg", no_complex_overload_reg);
+    m.def("complex_overload_reg", complex_overload_reg);
+
     m.def("readme_example1", readme_example1);
     m.def("readme_example2", xt::pyvectorize(readme_example2));
 

test_python/test_pyarray.py

@@ -36,6 +36,23 @@ def test_readme_example1(self):
         y = xt.readme_example1(v)
         np.testing.assert_allclose(y, 1.2853996391883833, 1e-12)
 
+    def test_complex_overload_reg(self):
+        a = 23.23
+        c = 2.0 + 3.1j
+        self.assertEqual(xt.complex_overload_reg(a), a)
+        self.assertEqual(xt.complex_overload_reg(c), c)
+
+    def test_complex_overload(self):
+        a = np.random.rand(3, 3)
+        b = np.random.rand(3, 3)
+        c = a + b * 1j
+        y = xt.complex_overload(c)
+        np.testing.assert_allclose(np.imag(y), np.imag(c))
+        np.testing.assert_allclose(np.real(y), np.real(c))
+        x = xt.complex_overload(b)
+        self.assertEqual(x.dtype, b.dtype)
+        np.testing.assert_allclose(x, b)
+
     def test_readme_example2(self):
         x = np.arange(15).reshape(3, 5)
         y = [1, 2, 3, 4, 5]