Add PyIter_NextItem() function (#113)

vstinner · web-flow · commit 38e2d327a31b · 2024-10-09T12:00:16.000+02:00
diff --git a/docs/api.rst b/docs/api.rst
@@ -33,6 +33,10 @@ Python 3.14
 
    See `PyLong_GetSign() documentation <https://docs.python.org/dev/c-api/long.html#c.PyLong_GetSign>`__.
 
+.. c:function:: PyObject* PyIter_NextItem(PyObject *sep, PyObject *iterable)
+
+   See `PyIter_NextItem() documentation <https://docs.python.org/dev/c-api/iter.html#c.PyIter_NextItem>`__.
+
 .. c:function:: PyObject* PyBytes_Join(PyObject *sep, PyObject *iterable)
 
    See `PyBytes_Join() documentation <https://docs.python.org/dev/c-api/bytes.html#c.PyBytes_Join>`__.
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -4,6 +4,7 @@ Changelog
 * 2024-10-09: Add functions:
 
   * ``PyBytes_Join()``
+  * ``PyIter_NextItem()``
   * ``PyUnicode_Equal()``
   * ``Py_HashBuffer()``
 
diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h
@@ -1519,14 +1519,12 @@ static inline int PyLong_GetSign(PyObject *obj, int *sign)
 static inline int PyUnicode_Equal(PyObject *str1, PyObject *str2)
 {
     if (!PyUnicode_Check(str1)) {
-        PyErr_Format(PyExc_TypeError,
-                     "first argument must be str, not %s",
+        PyErr_Format(PyExc_TypeError, "first argument must be str, not %s",
                      Py_TYPE(str1)->tp_name);
         return -1;
     }
     if (!PyUnicode_Check(str2)) {
-        PyErr_Format(PyExc_TypeError,
-                     "second argument must be str, not %s",
+        PyErr_Format(PyExc_TypeError, "second argument must be str, not %s",
                      Py_TYPE(str2)->tp_name);
         return -1;
     }
@@ -1576,6 +1574,37 @@ static inline Py_hash_t Py_HashBuffer(const void *ptr, Py_ssize_t len)
 #endif
 
 
+#if PY_VERSION_HEX < 0x030E00A0
+static inline int PyIter_NextItem(PyObject *iter, PyObject **item)
+{
+    iternextfunc tp_iternext;
+
+    assert(iter != NULL);
+    assert(item != NULL);
+
+    tp_iternext = Py_TYPE(iter)->tp_iternext;
+    if (tp_iternext == NULL) {
+        *item = NULL;
+        PyErr_Format(PyExc_TypeError, "expected an iterator, got '%s'",
+                     Py_TYPE(iter)->tp_name);
+        return -1;
+    }
+
+    if ((*item = tp_iternext(iter))) {
+        return 1;
+    }
+    if (!PyErr_Occurred()) {
+        return 0;
+    }
+    if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
+        PyErr_Clear();
+        return 0;
+    }
+    return -1;
+}
+#endif
+
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/tests/test_pythoncapi_compat_cext.c b/tests/test_pythoncapi_compat_cext.c
@@ -1437,12 +1437,14 @@ static int
 heapmanaged_traverse(PyObject *self, visitproc visit, void *arg)
 {
     Py_VISIT(Py_TYPE(self));
+    // Test PyObject_VisitManagedDict()
     return PyObject_VisitManagedDict(self, visit, arg);
 }
 
 static int
 heapmanaged_clear(PyObject *self)
 {
+    // Test PyObject_ClearManagedDict()
     PyObject_ClearManagedDict(self);
     return 0;
 }
@@ -1475,6 +1477,7 @@ static PyType_Spec HeapCTypeWithManagedDict_spec = {
 static PyObject *
 test_managed_dict(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 {
+    // Test PyObject_VisitManagedDict() and PyObject_ClearManagedDict()
     PyObject *type = PyType_FromSpec(&HeapCTypeWithManagedDict_spec);
     if (type == NULL) {
         return NULL;
@@ -1926,6 +1929,48 @@ test_bytes(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 }
 
 
+static PyObject *
+test_iter(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
+{
+    // Test PyIter_NextItem()
+    PyObject *tuple = Py_BuildValue("(i)", 123);
+    if (tuple == NULL) {
+        return NULL;
+    }
+    PyObject *iter = PyObject_GetIter(tuple);
+    Py_DECREF(tuple);
+    if (iter == NULL) {
+        return NULL;
+    }
+
+    // first item
+    PyObject *item = UNINITIALIZED_OBJ;
+    assert(PyIter_NextItem(iter, &item) == 1);
+    {
+        PyObject *expected = PyLong_FromLong(123);
+        assert(PyObject_RichCompareBool(item, expected, Py_EQ) == 1);
+        assert(expected != NULL);
+        Py_DECREF(expected);
+    }
+
+    // StopIteration
+    item = UNINITIALIZED_OBJ;
+    assert(PyIter_NextItem(iter, &item) == 0);
+    assert(item == NULL);
+    assert(!PyErr_Occurred());
+
+    // non-iterable object
+    item = UNINITIALIZED_OBJ;
+    assert(PyIter_NextItem(Py_None, &item) == -1);
+    assert(item == NULL);
+    assert(PyErr_ExceptionMatches(PyExc_TypeError));
+    PyErr_Clear();
+
+    Py_DECREF(iter);
+    Py_RETURN_NONE;
+}
+
+
 static struct PyMethodDef methods[] = {
     {"test_object", test_object, METH_NOARGS, _Py_NULL},
     {"test_py_is", test_py_is, METH_NOARGS, _Py_NULL},
@@ -1970,6 +2015,7 @@ static struct PyMethodDef methods[] = {
     {"test_unicodewriter_format", test_unicodewriter_format, METH_NOARGS, _Py_NULL},
 #endif
     {"test_bytes", test_bytes, METH_NOARGS, _Py_NULL},
+    {"test_iter", test_iter, METH_NOARGS, _Py_NULL},
     {_Py_NULL, _Py_NULL, 0, _Py_NULL}
 };
 
