Use nb::iterator instead of Py_Iterator

neuronsimulator · Oct 8, 2024 · a1fa94a · a1fa94a
1 parent ab0a571
commit a1fa94a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 11 deletions.
diff --git a/src/nrnpython/nrnpy_p2h.cpp b/src/nrnpython/nrnpy_p2h.cpp
@@ -14,6 +14,8 @@
 #include "parse.hpp"
 
 #include <nanobind/nanobind.h>
+
+#include "utils/enumerate.h"
 namespace nb = nanobind;
 
 static char* nrnpyerr_str();
@@ -937,15 +939,13 @@ static Object* py_alltoall_type(int size, int type) {
         // for alltoall, each rank handled identically
         // for scatter, root handled as list all, other ranks handled as None
         if (type == 1 || nrnmpi_myid == root) {  // psrc is list of nhost items
+            nb::list psrc_list(psrc);
 
             scnt = new int[np];
             for (int i = 0; i < np; ++i) {
                 scnt[i] = 0;
             }
 
-            PyObject* iterator = PyObject_GetIter(psrc);
-            PyObject* p;
-
             size_t bufsz = 100000;  // 100k buffer to start with
             if (size > 0) {         // or else the positive number specified
                 bufsz = size;
@@ -954,13 +954,13 @@ static Object* py_alltoall_type(int size, int type) {
                 s = new char[bufsz];
             }
             size_t curpos = 0;
-            for (size_t i = 0; (p = PyIter_Next(iterator)) != NULL; ++i) {
-                if (p == Py_None) {
+            for (auto&& [i, p]: enumerate(psrc_list)) {
+                if (p.is_none()) {
                     scnt[i] = 0;
-                    Py_DECREF(p);
+                    p.dec_ref();
                     continue;
                 }
-                auto b = pickle(p);
+                auto b = pickle(p.ptr());
                 if (size >= 0) {
                     if (curpos + b.size() >= bufsz) {
                         bufsz = bufsz * 2 + b.size();
@@ -977,9 +977,7 @@ static Object* py_alltoall_type(int size, int type) {
                 }
                 curpos += b.size();
                 scnt[i] = static_cast<int>(b.size());
-                Py_DECREF(p);
             }
-            Py_DECREF(iterator);
 
             // scatter equivalent to alltoall NONE list for not root ranks.
         } else if (type == 5 && nrnmpi_myid != root) {

diff --git a/src/utils/enumerate.h b/src/utils/enumerate.h
@@ -99,7 +99,7 @@ constexpr auto enumerate(T&& iterable) {
             ++iter;
         }
         auto operator*() const {
-            return std::tie(i, *iter);
+            return std::forward_as_tuple(i, *iter);
         }
     };
     struct iterable_wrapper {
@@ -129,7 +129,7 @@ constexpr auto renumerate(T&& iterable) {
             ++iter;
         }
         auto operator*() const {
-            return std::tie(i, *iter);
+            return std::forward_as_tuple(i, *iter);
         }
     };
     struct iterable_wrapper {