API: Make numpy.h compatible with both NumPy 1.x and 2.x

.github/workflows/ci.yml

@@ -108,12 +108,14 @@ jobs:
       run: python -m pip install pytest-github-actions-annotate-failures
 
     # First build - C++11 mode and inplace
-    # More-or-less randomly adding -DPYBIND11_SIMPLE_GIL_MANAGEMENT=ON here.
+    # More-or-less randomly adding -DPYBIND11_SIMPLE_GIL_MANAGEMENT=ON here
+    # (same for PYBIND11_NUMPY2_SUPPORT).
     - name: Configure C++11 ${{ matrix.args }}
       run: >
         cmake -S . -B .
         -DPYBIND11_WERROR=ON
         -DPYBIND11_SIMPLE_GIL_MANAGEMENT=ON
+        -DPYBIND11_NUMPY2_SUPPORT=ON
         -DDOWNLOAD_CATCH=ON
         -DDOWNLOAD_EIGEN=ON
         -DCMAKE_CXX_STANDARD=11
@@ -138,11 +140,13 @@ jobs:
 
     # Second build - C++17 mode and in a build directory
     # More-or-less randomly adding -DPYBIND11_SIMPLE_GIL_MANAGEMENT=OFF here.
+    # (same for PYBIND11_NUMPY2_SUPPORT, but requires a NumPy 1.x at runtime).
     - name: Configure C++17
       run: >
         cmake -S . -B build2
         -DPYBIND11_WERROR=ON
         -DPYBIND11_SIMPLE_GIL_MANAGEMENT=OFF
+        -DPYBIND11_NUMPY2_SUPPORT=OFF
         -DDOWNLOAD_CATCH=ON
         -DDOWNLOAD_EIGEN=ON
         -DCMAKE_CXX_STANDARD=17
@@ -895,8 +899,10 @@ jobs:
         python-version: ${{ matrix.python }}
 
     - name: Prepare env
+      # Ensure use of NumPy 2 (via NumPy nightlies but can be changed soon)
       run: |
         python3 -m pip install -r tests/requirements.txt
+        python3 -m pip install --upgrade --pre -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple numpy scipy
 
     - name: Update CMake
       uses: jwlawson/[email protected]
@@ -905,6 +911,7 @@ jobs:
       run: >
         cmake -S . -B build
         -DPYBIND11_WERROR=ON
+        -DPYBIND11_NUMPY2_SUPPORT=ON
         -DDOWNLOAD_CATCH=ON
         -DDOWNLOAD_EIGEN=ON
         -DCMAKE_CXX_STANDARD=20
@@ -925,6 +932,7 @@ jobs:
       run: >
         cmake -S . -B build_partial
         -DPYBIND11_WERROR=ON
+        -DPYBIND11_NUMPY2_SUPPORT=ON
         -DDOWNLOAD_CATCH=ON
         -DDOWNLOAD_EIGEN=ON
         -DCMAKE_CXX_STANDARD=20

CMakeLists.txt

@@ -109,13 +109,17 @@ option(PYBIND11_TEST "Build pybind11 test suite?" ${PYBIND11_MASTER_PROJECT})
 option(PYBIND11_NOPYTHON "Disable search for Python" OFF)
 option(PYBIND11_SIMPLE_GIL_MANAGEMENT
        "Use simpler GIL management logic that does not support disassociation" OFF)
+option(PYBIND11_NUMPY2_SUPPORT "Enable compile in a NumPy 2 and 1 compatible way" OFF)
 set(PYBIND11_INTERNALS_VERSION
     ""
     CACHE STRING "Override the ABI version, may be used to enable the unstable ABI.")
 
 if(PYBIND11_SIMPLE_GIL_MANAGEMENT)
   add_compile_definitions(PYBIND11_SIMPLE_GIL_MANAGEMENT)
 endif()
+if(PYBIND11_NUMPY2_SUPPORT)
+  add_compile_definitions(PYBIND11_NUMPY2_SUPPORT)
+endif()
 
 cmake_dependent_option(
   USE_PYTHON_INCLUDE_DIR

include/pybind11/numpy.h

@@ -32,7 +32,8 @@
 /* This will be true on all flat address space platforms and allows us to reduce the
    whole npy_intp / ssize_t / Py_intptr_t business down to just ssize_t for all size
    and dimension types (e.g. shape, strides, indexing), instead of inflicting this
-   upon the library user. */
+   upon the library user.
+   Note that NumPy 2 now uses ssize_t for `npy_intp` to simplify this. */
 static_assert(sizeof(::pybind11::ssize_t) == sizeof(Py_intptr_t), "ssize_t != Py_intptr_t");
 static_assert(std::is_signed<Py_intptr_t>::value, "Py_intptr_t must be signed");
 // We now can reinterpret_cast between py::ssize_t and Py_intptr_t (MSVC + PyPy cares)
@@ -53,7 +54,8 @@ struct handle_type_name<array> {
 template <typename type, typename SFINAE = void>
 struct npy_format_descriptor;
 
-struct PyArrayDescr_Proxy {
+/* NumPy 1 proxy (always includes legacy fields) */
+struct PyArrayDescr1_Proxy {
     PyObject_HEAD
     PyObject *typeobj;
     char kind;
@@ -68,6 +70,43 @@ struct PyArrayDescr_Proxy {
     PyObject *names;
 };
 
+#ifdef PYBIND11_NUMPY2_SUPPORT
+struct PyArrayDescr_Proxy {
+    PyObject_HEAD
+    PyObject *typeobj;
+    char kind;
+    char type;
+    char byteorder;
+    char _former_flags;
+    int type_num;
+    /* Additional fields are NumPy version specific. */
+};
+#else
+/* NumPy 1.x only, we can expose all fields */
+using PyArrayDescr_Proxy = PyArrayDescr1_Proxy;
+#endif
+
+/* NumPy 2 proxy, including legacy fields */
+struct PyArrayDescr2_Proxy {
+    PyObject_HEAD
+    PyObject *typeobj;
+    char kind;
+    char type;
+    char byteorder;
+    char _former_flags;
+    int type_num;
+    std::uint64_t flags;
+    ssize_t elsize;
+    ssize_t alignment;
+    PyObject *metadata;
+    Py_hash_t hash;
+    void *reserved_null[2];
+    /* The following fields only exist if 0 <= type_num < 2056 */
+    char *subarray;
+    PyObject *fields;
+    PyObject *names;
+};
+
 struct PyArray_Proxy {
     PyObject_HEAD
     char *data;
@@ -131,6 +170,13 @@ PYBIND11_NOINLINE module_ import_numpy_core_submodule(const char *submodule_name
     object numpy_version = numpy_lib.attr("NumpyVersion")(version_string);
     int major_version = numpy_version.attr("major").cast<int>();
 
+#ifndef PYBIND11_NUMPY2_SUPPORT
+    if (major_version >= 2) {
+        throw std::runtime_error("module compiled without NumPy 2 support. Please define "
+                                 "PYBIND11_NUMPY2_SUPPORT before including `numpy2.h` "
+                                 "or define it during build to enable NumPy 2 support.");
+    }
+#endif
     /* `numpy.core` was renamed to `numpy._core` in NumPy 2.0 as it officially
         became a private module. */
     std::string numpy_core_path = major_version >= 2 ? "numpy._core" : "numpy.core";
@@ -203,6 +249,8 @@ struct npy_api {
             NPY_ULONG_, NPY_ULONGLONG_, NPY_UINT_),
     };
 
+    unsigned int PyArray_RUNTIME_VERSION_;
+
     struct PyArray_Dims {
         Py_intptr_t *ptr;
         int len;
@@ -241,6 +289,7 @@ struct npy_api {
     PyObject *(*PyArray_FromAny_)(PyObject *, PyObject *, int, int, int, PyObject *);
     int (*PyArray_DescrConverter_)(PyObject *, PyObject **);
     bool (*PyArray_EquivTypes_)(PyObject *, PyObject *);
+#ifndef PYBIND11_NUMPY2_SUPPORT
     int (*PyArray_GetArrayParamsFromObject_)(PyObject *,
                                              PyObject *,
                                              unsigned char,
@@ -249,6 +298,7 @@ struct npy_api {
                                              Py_intptr_t *,
                                              PyObject **,
                                              PyObject *);
+#endif
     PyObject *(*PyArray_Squeeze_)(PyObject *);
     // Unused. Not removed because that affects ABI of the class.
     int (*PyArray_SetBaseObject_)(PyObject *, PyObject *);
@@ -275,7 +325,9 @@ struct npy_api {
         API_PyArray_View = 137,
         API_PyArray_DescrConverter = 174,
         API_PyArray_EquivTypes = 182,
+#ifndef PYBIND11_NUMPY2_SUPPORT
         API_PyArray_GetArrayParamsFromObject = 278,
+#endif
         API_PyArray_SetBaseObject = 282
     };
 
@@ -290,7 +342,8 @@ struct npy_api {
         npy_api api;
 #define DECL_NPY_API(Func) api.Func##_ = (decltype(api.Func##_)) api_ptr[API_##Func];
         DECL_NPY_API(PyArray_GetNDArrayCFeatureVersion);
-        if (api.PyArray_GetNDArrayCFeatureVersion_() < 0x7) {
+        api.PyArray_RUNTIME_VERSION_ = api.PyArray_GetNDArrayCFeatureVersion_();
+        if (api.PyArray_RUNTIME_VERSION_ < 0x7) {
             pybind11_fail("pybind11 numpy support requires numpy >= 1.7.0");
         }
         DECL_NPY_API(PyArray_Type);
@@ -309,7 +362,9 @@ struct npy_api {
         DECL_NPY_API(PyArray_View);
         DECL_NPY_API(PyArray_DescrConverter);
         DECL_NPY_API(PyArray_EquivTypes);
+#ifndef PYBIND11_NUMPY2_SUPPORT
         DECL_NPY_API(PyArray_GetArrayParamsFromObject);
+#endif
         DECL_NPY_API(PyArray_SetBaseObject);
 
 #undef DECL_NPY_API
@@ -331,6 +386,14 @@ inline const PyArrayDescr_Proxy *array_descriptor_proxy(const PyObject *ptr) {
     return reinterpret_cast<const PyArrayDescr_Proxy *>(ptr);
 }
 
+inline const PyArrayDescr1_Proxy *array_descriptor1_proxy(const PyObject *ptr) {
+    return reinterpret_cast<const PyArrayDescr1_Proxy *>(ptr);
+}
+
+inline const PyArrayDescr2_Proxy *array_descriptor2_proxy(const PyObject *ptr) {
+    return reinterpret_cast<const PyArrayDescr2_Proxy *>(ptr);
+}
+
 inline bool check_flags(const void *ptr, int flag) {
     return (flag == (array_proxy(ptr)->flags & flag));
 }
@@ -610,10 +673,31 @@ class dtype : public object {
     }
 
     /// Size of the data type in bytes.
+#ifndef PYBIND11_NUMPY2_SUPPORT
     ssize_t itemsize() const { return detail::array_descriptor_proxy(m_ptr)->elsize; }
+#else
+    ssize_t itemsize() const {
+        if (detail::npy_api::get().PyArray_RUNTIME_VERSION_ < 0x12) {
+            return detail::array_descriptor1_proxy(m_ptr)->elsize;
+        }
+        return detail::array_descriptor2_proxy(m_ptr)->elsize;
+    }
+#endif
 
     /// Returns true for structured data types.
+#ifndef PYBIND11_NUMPY2_SUPPORT
     bool has_fields() const { return detail::array_descriptor_proxy(m_ptr)->names != nullptr; }
+#else
+    bool has_fields() const {
+        if (detail::npy_api::get().PyArray_RUNTIME_VERSION_ < 0x12) {
+            return detail::array_descriptor1_proxy(m_ptr)->names != nullptr;
+        }
+        if (num() < 0 || num() >= 2056) {
+            return false;
+        }
+        return detail::array_descriptor2_proxy(m_ptr)->names != nullptr;
+    }
+#endif
 
     /// Single-character code for dtype's kind.
     /// For example, floating point types are 'f' and integral types are 'i'.
@@ -639,11 +723,29 @@ class dtype : public object {
     /// Single character for byteorder
     char byteorder() const { return detail::array_descriptor_proxy(m_ptr)->byteorder; }
 
-    /// Alignment of the data type
+/// Alignment of the data type
+#ifndef PYBIND11_NUMPY2_SUPPORT
     int alignment() const { return detail::array_descriptor_proxy(m_ptr)->alignment; }
+#else
+    ssize_t alignment() const {
+        if (detail::npy_api::get().PyArray_RUNTIME_VERSION_ < 0x12) {
+            return detail::array_descriptor1_proxy(m_ptr)->alignment;
+        }
+        return detail::array_descriptor2_proxy(m_ptr)->alignment;
+    }
+#endif
 
-    /// Flags for the array descriptor
+/// Flags for the array descriptor
+#ifndef PYBIND11_NUMPY2_SUPPORT
     char flags() const { return detail::array_descriptor_proxy(m_ptr)->flags; }
+#else
+    std::uint64_t flags() const {
+        if (detail::npy_api::get().PyArray_RUNTIME_VERSION_ < 0x12) {
+            return (unsigned char) detail::array_descriptor1_proxy(m_ptr)->flags;
+        }
+        return detail::array_descriptor2_proxy(m_ptr)->flags;
+    }
+#endif
 
 private:
     static object &_dtype_from_pep3118() {
@@ -810,9 +912,7 @@ class array : public buffer {
     }
 
     /// Byte size of a single element
-    ssize_t itemsize() const {
-        return detail::array_descriptor_proxy(detail::array_proxy(m_ptr)->descr)->elsize;
-    }
+    ssize_t itemsize() const { return dtype().itemsize(); }
 
     /// Total number of bytes
     ssize_t nbytes() const { return size() * itemsize(); }

tests/conftest.py

@@ -218,4 +218,5 @@ def pytest_report_header(config):
         f" {pybind11_tests.cpp_std}"
         f" {pybind11_tests.PYBIND11_INTERNALS_ID}"
         f" PYBIND11_SIMPLE_GIL_MANAGEMENT={pybind11_tests.PYBIND11_SIMPLE_GIL_MANAGEMENT}"
+        f" PYBIND11_NUMPY2_SUPPORT={pybind11_tests.PYBIND11_NUMPY2_SUPPORT}"
     )

tests/pybind11_tests.cpp

@@ -95,6 +95,12 @@ PYBIND11_MODULE(pybind11_tests, m) {
 #else
         false;
 #endif
+    m.attr("PYBIND11_NUMPY2_SUPPORT") =
+#if defined(PYBIND11_NUMPY2_SUPPORT)
+        true;
+#else
+        false;
+#endif
 
     bind_ConstructorStats(m);
 

tests/test_numpy_dtypes.py

@@ -3,6 +3,7 @@
 import pytest
 
 import env  # noqa: F401
+from pybind11_tests import PYBIND11_NUMPY2_SUPPORT
 from pybind11_tests import numpy_dtypes as m
 
 np = pytest.importorskip("numpy")
@@ -178,7 +179,12 @@ def test_dtype(simple_dtype):
     assert m.test_dtype_num() == [np.dtype(ch).num for ch in expected_chars]
     assert m.test_dtype_byteorder() == [np.dtype(ch).byteorder for ch in expected_chars]
     assert m.test_dtype_alignment() == [np.dtype(ch).alignment for ch in expected_chars]
-    assert m.test_dtype_flags() == [chr(np.dtype(ch).flags) for ch in expected_chars]
+    if PYBIND11_NUMPY2_SUPPORT:
+        assert m.test_dtype_flags() == [np.dtype(ch).flags for ch in expected_chars]
+    else:
+        assert m.test_dtype_flags() == [
+            chr(np.dtype(ch).flags) for ch in expected_chars
+        ]
 
 
 def test_recarray(simple_dtype, packed_dtype):