Merge pull request #1 from EricCousineau-TRI/issue/1328

numpy: Provide concrete size aliases, test equivalence for `dtype(...).num

Author: Cody Griffin

include/pybind11/numpy.h

@@ -14,6 +14,7 @@
 #include <numeric>
 #include <algorithm>
 #include <array>
+#include <cstdint>
 #include <cstdlib>
 #include <cstring>
 #include <sstream>
@@ -108,6 +109,18 @@ inline numpy_internals& get_numpy_internals() {
     return *ptr;
 }
 
+template <typename T> struct same_size {
+    template <typename U> using as = bool_constant<sizeof(T) == sizeof(U)>;
+};
+
+// Lookup a type according to its size, and return a value corresponding to the NumPy typenum.
+template <typename Concrete, typename... Check>
+constexpr int platform_lookup(const std::array<int, sizeof...(Check)> codes) {
+    using code_index = std::integral_constant<int, constexpr_first<same_size<Concrete>::template as, Check...>()>;
+    static_assert(code_index::value != sizeof...(Check), "Unable to match type on this platform");
+    return codes[code_index::value];
+}
+
 struct npy_api {
     enum constants {
         NPY_ARRAY_C_CONTIGUOUS_ = 0x0001,
@@ -126,7 +139,23 @@ struct npy_api {
         NPY_FLOAT_, NPY_DOUBLE_, NPY_LONGDOUBLE_,
         NPY_CFLOAT_, NPY_CDOUBLE_, NPY_CLONGDOUBLE_,
         NPY_OBJECT_ = 17,
-        NPY_STRING_, NPY_UNICODE_, NPY_VOID_
+        NPY_STRING_, NPY_UNICODE_, NPY_VOID_,
+        // Platform-dependent normalization
+        NPY_INT8_ = NPY_BYTE_,
+        NPY_UINT8_ = NPY_UBYTE_,
+        NPY_INT16_ = NPY_SHORT_,
+        NPY_UINT16_ = NPY_USHORT_,
+        // `npy_common.h` defines the integer aliases. In order, it checks:
+        // NPY_BITSOF_LONG, NPY_BITSOF_LONGLONG, NPY_BITSOF_INT, NPY_BITSOF_SHORT, NPY_BITSOF_CHAR
+        // and assigns the alias to the first matching size, so we should check in this order.
+        NPY_INT32_ = platform_lookup<std::int32_t, long, int, short>({{
+            NPY_LONG_, NPY_INT_, NPY_SHORT_}}),
+        NPY_UINT32_ = platform_lookup<std::uint32_t, unsigned long, unsigned int, unsigned short>({{
+            NPY_ULONG_, NPY_UINT_, NPY_USHORT_}}),
+        NPY_INT64_ = platform_lookup<std::int64_t, long, long long, int>({{
+            NPY_LONG_, NPY_LONGLONG_, NPY_INT_}}),
+        NPY_UINT64_ = platform_lookup<std::uint64_t, unsigned long, unsigned long long, unsigned int>({{
+            NPY_ULONG_, NPY_ULONGLONG_, NPY_UINT_}}),
     };
 
     typedef struct {
@@ -1004,8 +1033,8 @@ struct npy_format_descriptor<T, enable_if_t<satisfies_any_of<T, std::is_arithmet
     // NB: the order here must match the one in common.h
     constexpr static const int values[15] = {
         npy_api::NPY_BOOL_,
-        npy_api::NPY_BYTE_,   npy_api::NPY_UBYTE_,   npy_api::NPY_SHORT_,    npy_api::NPY_USHORT_,
-        npy_api::NPY_INT_,    npy_api::NPY_UINT_,    npy_api::NPY_LONGLONG_, npy_api::NPY_ULONGLONG_,
+        npy_api::NPY_BYTE_,   npy_api::NPY_UBYTE_,   npy_api::NPY_INT16_,    npy_api::NPY_UINT16_,
+        npy_api::NPY_INT32_,  npy_api::NPY_UINT32_,  npy_api::NPY_INT64_,    npy_api::NPY_UINT64_,
         npy_api::NPY_FLOAT_,  npy_api::NPY_DOUBLE_,  npy_api::NPY_LONGDOUBLE_,
         npy_api::NPY_CFLOAT_, npy_api::NPY_CDOUBLE_, npy_api::NPY_CLONGDOUBLE_
     };

tests/test_numpy_array.cpp

@@ -14,6 +14,67 @@
 
 #include <cstdint>
 
+// Size / dtype checks.
+struct DtypeCheck {
+    py::dtype numpy{};
+    py::dtype pybind11{};
+};
+
+template <typename T>
+DtypeCheck get_dtype_check(const char* name) {
+    py::module np = py::module::import("numpy");
+    DtypeCheck check{};
+    check.numpy = np.attr("dtype")(np.attr(name));
+    check.pybind11 = py::dtype::of<T>();
+    return check;
+}
+
+std::vector<DtypeCheck> get_concrete_dtype_checks() {
+    return {
+        // Normalization
+        get_dtype_check<std::int8_t>("int8"),
+        get_dtype_check<std::uint8_t>("uint8"),
+        get_dtype_check<std::int16_t>("int16"),
+        get_dtype_check<std::uint16_t>("uint16"),
+        get_dtype_check<std::int32_t>("int32"),
+        get_dtype_check<std::uint32_t>("uint32"),
+        get_dtype_check<std::int64_t>("int64"),
+        get_dtype_check<std::uint64_t>("uint64")
+    };
+}
+
+struct DtypeSizeCheck {
+    std::string name{};
+    int size_cpp{};
+    int size_numpy{};
+    // For debugging.
+    py::dtype dtype{};
+};
+
+template <typename T>
+DtypeSizeCheck get_dtype_size_check() {
+    DtypeSizeCheck check{};
+    check.name = py::type_id<T>();
+    check.size_cpp = sizeof(T);
+    check.dtype = py::dtype::of<T>();
+    check.size_numpy = check.dtype.attr("itemsize").template cast<int>();
+    return check;
+}
+
+std::vector<DtypeSizeCheck> get_platform_dtype_size_checks() {
+    return {
+        get_dtype_size_check<short>(),
+        get_dtype_size_check<unsigned short>(),
+        get_dtype_size_check<int>(),
+        get_dtype_size_check<unsigned int>(),
+        get_dtype_size_check<long>(),
+        get_dtype_size_check<unsigned long>(),
+        get_dtype_size_check<long long>(),
+        get_dtype_size_check<unsigned long long>(),
+    };
+}
+
+// Arrays.
 using arr = py::array;
 using arr_t = py::array_t<uint16_t, 0>;
 static_assert(std::is_same<arr_t::value_type, uint16_t>::value, "");
@@ -72,6 +133,26 @@ TEST_SUBMODULE(numpy_array, sm) {
     try { py::module::import("numpy"); }
     catch (...) { return; }
 
+    // test_dtypes
+    py::class_<DtypeCheck>(sm, "DtypeCheck")
+        .def_readonly("numpy", &DtypeCheck::numpy)
+        .def_readonly("pybind11", &DtypeCheck::pybind11)
+        .def("__repr__", [](const DtypeCheck& self) {
+            return py::str("<DtypeCheck numpy={} pybind11={}>").format(
+                self.numpy, self.pybind11);
+        });
+    sm.def("get_concrete_dtype_checks", &get_concrete_dtype_checks);
+
+    py::class_<DtypeSizeCheck>(sm, "DtypeSizeCheck")
+        .def_readonly("name", &DtypeSizeCheck::name)
+        .def_readonly("size_cpp", &DtypeSizeCheck::size_cpp)
+        .def_readonly("size_numpy", &DtypeSizeCheck::size_numpy)
+        .def("__repr__", [](const DtypeSizeCheck& self) {
+            return py::str("<DtypeSizeCheck name='{}' size_cpp={} size_numpy={} dtype={}>").format(
+                self.name, self.size_cpp, self.size_numpy, self.dtype);
+        });
+    sm.def("get_platform_dtype_size_checks", &get_platform_dtype_size_checks);
+
     // test_array_attributes
     sm.def("ndim", [](const arr& a) { return a.ndim(); });
     sm.def("shape", [](const arr& a) { return arr(a.ndim(), a.shape()); });

tests/test_numpy_array.py

@@ -7,6 +7,21 @@
     import numpy as np
 
 
+def test_dtypes():
+    # See issue #1328.
+    # - Platform-dependent sizes.
+    for size_check in m.get_platform_dtype_size_checks():
+        print(size_check)
+        assert size_check.size_cpp == size_check.size_numpy, size_check
+    # - Concrete sizes.
+    for check in m.get_concrete_dtype_checks():
+        print(check)
+        assert check.numpy == check.pybind11, check
+        if check.numpy.num != check.pybind11.num:
+            print("NOTE: typenum mismatch for {}: {} != {}".format(
+                check, check.numpy.num, check.pybind11.num))
+
+
 @pytest.fixture(scope='function')
 def arr():
     return np.array([[1, 2, 3], [4, 5, 6]], '=u2')