support char arrays and complicated dtypes (structs) (xtensor-stack#149)

wolfv · SylvainCorlay · commit 1d6429f14798 · 2018-06-20T21:31:07.000+02:00
* add support for char arrays and structs

* ..
diff --git a/include/xtensor-python/pyarray.hpp b/include/xtensor-python/pyarray.hpp
@@ -49,12 +49,7 @@ namespace pybind11
             {
                 if (!convert)
                 {
-                    if (!PyArray_Check(src.ptr()))
-                    {
-                        return false;
-                    }
-                    int type_num = xt::detail::numpy_traits<T>::type_num;
-                    if(xt::detail::pyarray_type(reinterpret_cast<PyArrayObject*>(src.ptr())) != type_num)
+                    if (!xt::detail::check_array<T>(src))
                     {
                         return false;
                     }
@@ -477,7 +472,7 @@ namespace xt
         shape_type shape = xtl::make_sequence<shape_type>(0, size_type(1));
         strides_type strides = xtl::make_sequence<strides_type>(0, size_type(0));
         init_array(shape, strides);
-        m_storage[0] = T();
+        detail::default_initialize(m_storage);
     }
 
     /**
@@ -705,13 +700,15 @@ namespace xt
         {
             flags |= NPY_ARRAY_WRITEABLE;
         }
-        int type_num = detail::numpy_traits<T>::type_num;
+
+        auto dtype = pybind11::detail::npy_format_descriptor<T>::dtype();
 
         npy_intp* shape_data = reinterpret_cast<npy_intp*>(const_cast<size_type*>(shape.data()));
         npy_intp* strides_data = reinterpret_cast<npy_intp*>(adapted_strides.data());
+
         auto tmp = pybind11::reinterpret_steal<pybind11::object>(
-            PyArray_New(&PyArray_Type, static_cast<int>(shape.size()), shape_data, type_num, strides_data,
-                        nullptr, static_cast<int>(sizeof(T)), flags, nullptr));
+            PyArray_NewFromDescr(&PyArray_Type, (PyArray_Descr*) dtype.release().ptr(), static_cast<int>(shape.size()), shape_data, strides_data,
+                        nullptr, flags, nullptr));
 
         if (!tmp)
         {
diff --git a/include/xtensor-python/pycontainer.hpp b/include/xtensor-python/pycontainer.hpp
@@ -15,6 +15,7 @@
 
 #include "pybind11/complex.h"
 #include "pybind11/pybind11.h"
+#include "pybind11/numpy.h"
 
 #ifndef FORCE_IMPORT_ARRAY
 #define NO_IMPORT_ARRAY
@@ -129,8 +130,11 @@ namespace xt
 
     namespace detail
     {
+        template <class T, class E = void>
+        struct numpy_traits;
+
         template <class T>
-        struct numpy_traits
+        struct numpy_traits<T, std::enable_if_t<pybind11::detail::satisfies_any_of<T, std::is_arithmetic, xtl::is_complex>::value>>
         {
         private:
 
@@ -184,6 +188,47 @@ namespace xt
         {
             return numpy_enum_adjuster<NPY_LONGLONG != NPY_INT64>::pyarray_type(obj);
         }
+
+        template <class T>
+        void default_initialize_impl(T& storage, std::false_type)
+        {
+        }
+
+        template <class T>
+        void default_initialize_impl(T& storage, std::true_type)
+        {
+            using value_type = typename T::value_type;
+            storage[0] = value_type{};
+        }
+
+        template <class T>
+        void default_initialize(T& storage)
+        {
+            using value_type = typename T::value_type;
+            default_initialize_impl(storage, std::is_copy_assignable<value_type>());
+        }
+
+        template <class T>
+        bool check_array_type(const pybind11::handle& src, std::true_type)
+        {
+            int type_num = xt::detail::numpy_traits<T>::type_num;
+            return xt::detail::pyarray_type(reinterpret_cast<PyArrayObject*>(src.ptr())) == type_num;
+        }
+
+        template <class T>
+        bool check_array_type(const pybind11::handle& src, std::false_type)
+        {
+            return PyArray_EquivTypes((PyArray_Descr*) pybind11::detail::array_proxy(src.ptr())->descr,
+                                      (PyArray_Descr*) pybind11::dtype::of<T>().ptr());
+        }
+
+        template <class T>
+        bool check_array(const pybind11::handle& src)
+        {
+            using is_arithmetic_type = std::integral_constant<bool, bool(pybind11::detail::satisfies_any_of<T, std::is_arithmetic, xtl::is_complex>::value)>;
+            return PyArray_Check(src.ptr()) &&
+                   check_array_type<T>(src, is_arithmetic_type{});
+        }
     }
 
     /******************************
@@ -232,9 +277,9 @@ namespace xt
     template <class D>
     inline bool pycontainer<D>::check_(pybind11::handle h)
     {
-        int type_num = detail::numpy_traits<value_type>::type_num;
+        auto dtype = pybind11::detail::npy_format_descriptor<value_type>::dtype();
         return PyArray_Check(h.ptr()) &&
-            PyArray_EquivTypenums(PyArray_TYPE(reinterpret_cast<PyArrayObject*>(h.ptr())), type_num);
+            PyArray_EquivTypes_(PyArray_TYPE(reinterpret_cast<PyArrayObject*>(h.ptr())), dtype.ptr());
     }
 
     template <class D>
@@ -244,8 +289,9 @@ namespace xt
         {
             return nullptr;
         }
-        int type_num = detail::numpy_traits<value_type>::type_num;
-        auto res = PyArray_FromAny(ptr, PyArray_DescrFromType(type_num), 0, 0,
+
+        auto dtype = pybind11::detail::npy_format_descriptor<value_type>::dtype();
+        auto res = PyArray_FromAny(ptr, (PyArray_Descr *) dtype.release().ptr(), 0, 0,
                                    NPY_ARRAY_ENSUREARRAY | NPY_ARRAY_FORCECAST, nullptr);
         return res;
     }
diff --git a/include/xtensor-python/pytensor.hpp b/include/xtensor-python/pytensor.hpp
@@ -50,12 +50,7 @@ namespace pybind11
             {
                 if (!convert)
                 {
-                    if (!PyArray_Check(src.ptr()))
-                    {
-                        return false;
-                    }
-                    int type_num = xt::detail::numpy_traits<T>::type_num;
-                    if(xt::detail::pyarray_type(reinterpret_cast<PyArrayObject*>(src.ptr())) != type_num)
+                    if (!xt::detail::check_array<T>(src))
                     {
                         return false;
                     }
@@ -228,7 +223,7 @@ namespace xt
         m_shape = xtl::make_sequence<shape_type>(N, size_type(1));
         m_strides = xtl::make_sequence<strides_type>(N, size_type(0));
         init_tensor(m_shape, m_strides);
-        m_storage[0] = T();
+        detail::default_initialize(m_storage);
     }
 
     /**
@@ -402,11 +397,12 @@ namespace xt
         {
             flags |= NPY_ARRAY_WRITEABLE;
         }
-        int type_num = detail::numpy_traits<T>::type_num;
+        auto dtype = pybind11::detail::npy_format_descriptor<T>::dtype();
 
         auto tmp = pybind11::reinterpret_steal<pybind11::object>(
-            PyArray_New(&PyArray_Type, N, const_cast<npy_intp*>(shape.data()),
-                        type_num, python_strides, nullptr, sizeof(T), flags, nullptr));
+            PyArray_NewFromDescr(&PyArray_Type, (PyArray_Descr*) dtype.ptr(), static_cast<int>(shape.size()),
+                        const_cast<npy_intp*>(shape.data()), python_strides,
+                        nullptr, flags, nullptr));
 
         if (!tmp)
         {
diff --git a/test_python/main.cpp b/test_python/main.cpp
@@ -107,6 +107,53 @@ void dump_numpy_constant()
     std::cout << "NPY_UINT64 = " << NPY_UINT64 << std::endl;
 }
 
+struct A
+{
+    double a;
+    int b;
+    char c;
+    std::array<double, 3> x;
+};
+
+struct B
+{
+    double a;
+    int b;
+};
+
+xt::pyarray<A> dtype_to_python()
+{
+    A a1{123, 321, 'a', {1, 2, 3}};
+    A a2{111, 222, 'x', {5, 5, 5}};
+
+    return xt::pyarray<A>({a1, a2});
+}
+
+xt::pyarray<B> dtype_from_python(xt::pyarray<B>& b)
+{
+    if (b(0).a != 1 || b(0).b != 'p' || b(1).a != 123 || b(1).b != 'c')
+    {
+        throw std::runtime_error("FAIL");
+    }
+
+    b(0).a = 123.;
+    b(0).b = 'w';
+    return b;
+}
+
+void char_array(xt::pyarray<char[20]>& carr)
+{
+    if (strcmp(carr(2), "python"))
+    {
+        throw std::runtime_error("TEST FAILED!");
+    }
+    std::fill(&carr(2)[0], &carr(2)[0] + 20, 0);
+    carr(2)[0] = 'c';
+    carr(2)[1] = '+';
+    carr(2)[2] = '+';
+    carr(2)[3] = '\0';
+}
+
 PYBIND11_MODULE(xtensor_python_test, m)
 {
     xt::import_numpy();
@@ -142,4 +189,12 @@ PYBIND11_MODULE(xtensor_python_test, m)
     m.def("int_overload", int_overload<int64_t>);
 
     m.def("dump_numpy_constant", dump_numpy_constant);
+
+    // Register additional dtypes
+    PYBIND11_NUMPY_DTYPE(A, a, b, c, x);
+    PYBIND11_NUMPY_DTYPE(B, a, b);
+
+    m.def("dtype_to_python", dtype_to_python);
+    m.def("dtype_from_python", dtype_from_python);
+    m.def("char_array", char_array);
 }
diff --git a/test_python/test_pyarray.py b/test_python/test_pyarray.py
@@ -21,7 +21,7 @@
 import xtensor_python_test as xt
 import numpy as np
 
-class ExampleTest(TestCase):
+class XtensorTest(TestCase):
 
     def test_example1(self):
         self.assertEqual(4, xt.example1([4, 5, 6]))
@@ -87,3 +87,41 @@ def test_int_overload(self):
             b = xt.int_overload(np.ones((10), dtype))
             self.assertEqual(str(dtype.__name__), b)
 
+    def test_dtype(self):
+        var = xt.dtype_to_python()
+        self.assertEqual(var.dtype.names, ('a', 'b', 'c', 'x'))
+
+        exp_dtype = {
+             'a': (np.dtype('float64'), 0),
+             'b': (np.dtype('int32'), 8),
+             'c': (np.dtype('int8'), 12),
+             'x': (np.dtype(('<f8', (3,))), 16)
+        }
+
+        self.assertEqual(var.dtype.fields, exp_dtype)
+
+        self.assertEqual(var[0]['a'], 123)
+        self.assertEqual(var[0]['b'], 321)
+        self.assertEqual(var[0]['c'], ord('a'))
+        self.assertTrue(np.all(var[0]['x'] == [1, 2, 3]))
+
+        self.assertEqual(var[1]['a'], 111)
+        self.assertEqual(var[1]['b'], 222)
+        self.assertEqual(var[1]['c'], ord('x'))
+        self.assertTrue(np.all(var[1]['x'] == [5, 5, 5]))
+
+        d_dtype = np.dtype({'names':['a','b'], 'formats':['<f8','<i4'], 'offsets':[0,8], 'itemsize':16})
+
+        darr = np.array([(1, ord('p')), (123, ord('c'))], dtype=d_dtype)
+        self.assertEqual(darr[0]['a'], 1)
+        res = xt.dtype_from_python(darr)
+        self.assertEqual(res[0]['a'], 123.)
+        self.assertEqual(darr[0]['a'], 123.)
+
+    def test_char_array(self):
+        var = np.array(['hello', 'from', 'python'], dtype=np.dtype('|S20'));
+        self.assertEqual(var[0], b'hello')
+        xt.char_array(var)
+        self.assertEqual(var[0], b'hello')
+        self.assertEqual(var[1], b'from')
+        self.assertEqual(var[2], b'c++')

Original file line number	Diff line number	Diff line change
`@@ -49,12 +49,7 @@ namespace pybind11`
`49`	`49`	`{`
`50`	`50`	`if (!convert)`
`51`	`51`	`{`
`52`		`- if (!PyArray_Check(src.ptr()))`
`53`		`- {`
`54`		`- return false;`
`55`		`- }`
`56`		`- int type_num = xt::detail::numpy_traits<T>::type_num;`
`57`		`- if(xt::detail::pyarray_type(reinterpret_cast<PyArrayObject*>(src.ptr())) != type_num)`
	`52`	`+ if (!xt::detail::check_array<T>(src))`
`58`	`53`	`{`
`59`	`54`	`return false;`
`60`	`55`	`}`
`@@ -477,7 +472,7 @@ namespace xt`
`477`	`472`	`shape_type shape = xtl::make_sequence<shape_type>(0, size_type(1));`
`478`	`473`	`strides_type strides = xtl::make_sequence<strides_type>(0, size_type(0));`
`479`	`474`	`init_array(shape, strides);`
`480`		`- m_storage[0] = T();`
	`475`	`+ detail::default_initialize(m_storage);`
`481`	`476`	`}`
`482`	`477`
`483`	`478`	`/**`
`@@ -705,13 +700,15 @@ namespace xt`
`705`	`700`	`{`
`706`	`701`	`flags \|= NPY_ARRAY_WRITEABLE;`
`707`	`702`	`}`
`708`		`- int type_num = detail::numpy_traits<T>::type_num;`
	`703`	`+`
	`704`	`+ auto dtype = pybind11::detail::npy_format_descriptor<T>::dtype();`
`709`	`705`
`710`	`706`	`npy_intp* shape_data = reinterpret_cast<npy_intp>(const_cast<size_type>(shape.data()));`
`711`	`707`	`npy_intp* strides_data = reinterpret_cast<npy_intp*>(adapted_strides.data());`
	`708`	`+`
`712`	`709`	`auto tmp = pybind11::reinterpret_steal<pybind11::object>(`
`713`		`- PyArray_New(&PyArray_Type, static_cast<int>(shape.size()), shape_data, type_num, strides_data,`
`714`		`- nullptr, static_cast<int>(sizeof(T)), flags, nullptr));`
	`710`	`+ PyArray_NewFromDescr(&PyArray_Type, (PyArray_Descr*) dtype.release().ptr(), static_cast<int>(shape.size()), shape_data, strides_data,`
	`711`	`+ nullptr, flags, nullptr));`
`715`	`712`
`716`	`713`	`if (!tmp)`
`717`	`714`	`{`
Original file line number	Diff line number	Diff line change
`@@ -50,12 +50,7 @@ namespace pybind11`
`50`	`50`	`{`
`51`	`51`	`if (!convert)`
`52`	`52`	`{`
`53`		`- if (!PyArray_Check(src.ptr()))`
`54`		`- {`
`55`		`- return false;`
`56`		`- }`
`57`		`- int type_num = xt::detail::numpy_traits<T>::type_num;`
`58`		`- if(xt::detail::pyarray_type(reinterpret_cast<PyArrayObject*>(src.ptr())) != type_num)`
	`53`	`+ if (!xt::detail::check_array<T>(src))`
`59`	`54`	`{`
`60`	`55`	`return false;`
`61`	`56`	`}`
`@@ -228,7 +223,7 @@ namespace xt`
`228`	`223`	`m_shape = xtl::make_sequence<shape_type>(N, size_type(1));`
`229`	`224`	`m_strides = xtl::make_sequence<strides_type>(N, size_type(0));`
`230`	`225`	`init_tensor(m_shape, m_strides);`
`231`		`- m_storage[0] = T();`
	`226`	`+ detail::default_initialize(m_storage);`
`232`	`227`	`}`
`233`	`228`
`234`	`229`	`/**`
`@@ -402,11 +397,12 @@ namespace xt`
`402`	`397`	`{`
`403`	`398`	`flags \|= NPY_ARRAY_WRITEABLE;`
`404`	`399`	`}`
`405`		`- int type_num = detail::numpy_traits<T>::type_num;`
	`400`	`+ auto dtype = pybind11::detail::npy_format_descriptor<T>::dtype();`
`406`	`401`
`407`	`402`	`auto tmp = pybind11::reinterpret_steal<pybind11::object>(`
`408`		`- PyArray_New(&PyArray_Type, N, const_cast<npy_intp*>(shape.data()),`
`409`		`- type_num, python_strides, nullptr, sizeof(T), flags, nullptr));`
	`403`	`+ PyArray_NewFromDescr(&PyArray_Type, (PyArray_Descr*) dtype.ptr(), static_cast<int>(shape.size()),`
	`404`	`+ const_cast<npy_intp*>(shape.data()), python_strides,`
	`405`	`+ nullptr, flags, nullptr));`
`410`	`406`
`411`	`407`	`if (!tmp)`
`412`	`408`	`{`