Convert _path.is_sorted_and_has_non_nan to pybind11

QuLogic · QuLogic · commit 1a9b60110b3c · 2024-09-23T19:24:17.000-04:00
diff --git a/src/_path.h b/src/_path.h
@@ -1223,17 +1223,15 @@ bool convert_to_string(PathIterator &path,
 }
 
 template<class T>
-bool is_sorted_and_has_non_nan(PyArrayObject *array)
+bool is_sorted_and_has_non_nan(py::array_t<T> array)
 {
-    char* ptr = PyArray_BYTES(array);
-    npy_intp size = PyArray_DIM(array, 0),
-             stride = PyArray_STRIDE(array, 0);
+    auto size = array.shape(0);
     using limits = std::numeric_limits<T>;
     T last = limits::has_infinity ? -limits::infinity() : limits::min();
     bool found_non_nan = false;
 
-    for (npy_intp i = 0; i < size; ++i, ptr += stride) {
-        T current = *(T*)ptr;
+    for (auto i = 0; i < size; ++i) {
+        T current = *array.data(i);
         // The following tests !isnan(current), but also works for integral
         // types.  (The isnan(IntegralType) overload is absent on MSVC.)
         if (current == current) {
diff --git a/src/_path_wrapper.cpp b/src/_path_wrapper.cpp
@@ -350,41 +350,26 @@ Py_is_sorted_and_has_non_nan(py::object obj)
 {
     bool result;
 
-    PyArrayObject *array = (PyArrayObject *)PyArray_CheckFromAny(
-        obj.ptr(), NULL, 1, 1, NPY_ARRAY_NOTSWAPPED, NULL);
-
-    if (array == NULL) {
-        throw py::error_already_set();
+    py::array array = py::array::ensure(obj);
+    if (array.ndim() != 1) {
+        throw std::invalid_argument("array must be 1D");
     }
 
+    auto dtype = array.dtype();
     /* Handle just the most common types here, otherwise coerce to double */
-    switch (PyArray_TYPE(array)) {
-    case NPY_INT:
-        result = is_sorted_and_has_non_nan<npy_int>(array);
-        break;
-    case NPY_LONG:
-        result = is_sorted_and_has_non_nan<npy_long>(array);
-        break;
-    case NPY_LONGLONG:
-        result = is_sorted_and_has_non_nan<npy_longlong>(array);
-        break;
-    case NPY_FLOAT:
-        result = is_sorted_and_has_non_nan<npy_float>(array);
-        break;
-    case NPY_DOUBLE:
-        result = is_sorted_and_has_non_nan<npy_double>(array);
-        break;
-    default:
-        Py_DECREF(array);
-        array = (PyArrayObject *)PyArray_FromObject(obj.ptr(), NPY_DOUBLE, 1, 1);
-        if (array == NULL) {
-            throw py::error_already_set();
-        }
-        result = is_sorted_and_has_non_nan<npy_double>(array);
+    if (dtype.equal(py::dtype::of<std::int32_t>())) {
+        result = is_sorted_and_has_non_nan<int32_t>(array);
+    } else if (dtype.equal(py::dtype::of<std::int64_t>())) {
+        result = is_sorted_and_has_non_nan<int64_t>(array);
+    } else if (dtype.equal(py::dtype::of<float>())) {
+        result = is_sorted_and_has_non_nan<float>(array);
+    } else if (dtype.equal(py::dtype::of<double>())) {
+        result = is_sorted_and_has_non_nan<double>(array);
+    } else {
+        array = py::array_t<double>::ensure(obj);
+        result = is_sorted_and_has_non_nan<double>(array);
     }
 
-    Py_DECREF(array);
-
     return result;
 }
 
