Add PyDict_Pop() function (#81)

vstinner · web-flow · commit 5bf2fb2f2d12 · 2023-11-14T15:02:42.000+01:00
Add PyDict_Pop() and PyDict_PopString() functions.
diff --git a/docs/api.rst b/docs/api.rst
@@ -129,6 +129,14 @@ Python 3.13
 
    See `PyList_Clear() documentation <https://docs.python.org/dev/c-api/list.html#c.PyList_Clear>`__.
 
+.. c:function:: int PyDict_Pop(PyObject *dict, PyObject *key, PyObject **result)
+
+   See `PyDict_Pop() documentation <https://docs.python.org/dev/c-api/dict.html#c.PyDict_Pop>`__.
+
+.. c:function:: int PyDict_PopString(PyObject *dict, const char *key, PyObject **result)
+
+   See `PyDict_PopString() documentation <https://docs.python.org/dev/c-api/dict.html#c.PyDict_PopString>`__.
+
 
 Python 3.12
 -----------
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -1,6 +1,11 @@
 Changelog
 =========
 
+* 2023-11-14: Add functions:
+
+  * ``PyDict_Pop()``
+  * ``PyDict_PopString()``
+
 * 2023-11-13: Add functions:
 
   * ``PyList_Extend()``
diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h
@@ -1028,6 +1028,67 @@ PyList_Clear(PyObject *list)
 }
 #endif
 
+// gh-111262 added PyDict_Pop() and PyDict_PopString() to Python 3.13.0a2
+#if PY_VERSION_HEX < 0x030D00A2
+static inline int
+PyDict_Pop(PyObject *dict, PyObject *key, PyObject **result)
+{
+    PyObject *value;
+
+    if (!PyDict_Check(dict)) {
+        PyErr_BadInternalCall();
+        if (result) {
+            *result = NULL;
+        }
+        return -1;
+    }
+
+    // bpo-16991 added _PyDict_Pop() to Python 3.5.0b2.
+    // Python 3.6.0b3 changed _PyDict_Pop() first argument type to PyObject*.
+    // Python 3.13.0a1 removed _PyDict_Pop().
+#if defined(PYPY_VERSION) || PY_VERSION_HEX < 0x030500b2 || PY_VERSION_HEX >= 0x030D0000
+    value = PyObject_CallMethod(dict, "pop", "O", key);
+#elif PY_VERSION_HEX < 0x030600b3
+    value = _PyDict_Pop(_Py_CAST(PyDictObject*, dict), key, NULL);
+#else
+    value = _PyDict_Pop(dict, key, NULL);
+#endif
+    if (value == NULL) {
+        if (result) {
+            *result = NULL;
+        }
+        if (PyErr_Occurred() && !PyErr_ExceptionMatches(PyExc_KeyError)) {
+            return -1;
+        }
+        PyErr_Clear();
+        return 0;
+    }
+    if (result) {
+        *result = value;
+    }
+    else {
+        Py_DECREF(value);
+    }
+    return 1;
+}
+
+static inline int
+PyDict_PopString(PyObject *dict, const char *key, PyObject **result)
+{
+    PyObject *key_obj = PyUnicode_FromString(key);
+    if (key_obj == NULL) {
+        if (result != NULL) {
+            *result = NULL;
+        }
+        return -1;
+    }
+
+    int res = PyDict_Pop(dict, key_obj, result);
+    Py_DECREF(key_obj);
+    return res;
+}
+#endif
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/tests/test_pythoncapi_compat_cext.c b/tests/test_pythoncapi_compat_cext.c
@@ -1253,6 +1253,82 @@ test_dict_api(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 }
 
 
+static PyObject *
+test_dict_pop(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
+{
+    PyObject *dict = PyDict_New();
+    if (dict == NULL) {
+        return NULL;
+    }
+
+    PyObject *key = PyUnicode_FromString("key");
+    assert(key != NULL);
+    PyObject *value = PyUnicode_FromString("abc");
+    assert(value != NULL);
+
+    // test PyDict_Pop(), get the removed value, key is present
+    assert(PyDict_SetItem(dict, key, value) == 0);
+    PyObject *removed = UNINITIALIZED_OBJ;
+    assert(PyDict_Pop(dict, key, &removed) == 1);
+    assert(removed == value);
+    Py_DECREF(removed);
+
+    // test PyDict_Pop(), ignore the removed value, key is present
+    assert(PyDict_SetItem(dict, key, value) == 0);
+    assert(PyDict_Pop(dict, key, NULL) == 1);
+
+    // test PyDict_Pop(), key is missing
+    removed = UNINITIALIZED_OBJ;
+    assert(PyDict_Pop(dict, key, &removed) == 0);
+    assert(removed == NULL);
+    assert(PyDict_Pop(dict, key, NULL) == 0);
+
+    // test PyDict_PopString(), get the removed value, key is present
+    assert(PyDict_SetItem(dict, key, value) == 0);
+    removed = UNINITIALIZED_OBJ;
+    assert(PyDict_PopString(dict, "key", &removed) == 1);
+    assert(removed == value);
+    Py_DECREF(removed);
+
+    // test PyDict_PopString(), ignore the removed value, key is present
+    assert(PyDict_SetItem(dict, key, value) == 0);
+    assert(PyDict_PopString(dict, "key", NULL) == 1);
+
+    // test PyDict_PopString(), key is missing
+    removed = UNINITIALIZED_OBJ;
+    assert(PyDict_PopString(dict, "key", &removed) == 0);
+    assert(removed == NULL);
+    assert(PyDict_PopString(dict, "key", NULL) == 0);
+
+    // dict error
+    removed = UNINITIALIZED_OBJ;
+    assert(PyDict_Pop(key, key, &removed) == -1);
+    assert(removed == NULL);
+    assert(PyErr_ExceptionMatches(PyExc_SystemError));
+    PyErr_Clear();
+
+    assert(PyDict_Pop(key, key, NULL) == -1);
+    assert(PyErr_ExceptionMatches(PyExc_SystemError));
+    PyErr_Clear();
+
+    removed = UNINITIALIZED_OBJ;
+    assert(PyDict_PopString(key, "key", &removed) == -1);
+    assert(removed == NULL);
+    assert(PyErr_ExceptionMatches(PyExc_SystemError));
+    PyErr_Clear();
+
+    assert(PyDict_PopString(key, "key", NULL) == -1);
+    assert(PyErr_ExceptionMatches(PyExc_SystemError));
+    PyErr_Clear();
+
+    // exit
+    Py_DECREF(dict);
+    Py_DECREF(key);
+    Py_DECREF(value);
+    Py_RETURN_NONE;
+}
+
+
 static PyObject *
 test_long_api(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 {
@@ -1447,6 +1523,7 @@ static struct PyMethodDef methods[] = {
     {"test_getattr", test_getattr, METH_NOARGS, _Py_NULL},
     {"test_getitem", test_getitem, METH_NOARGS, _Py_NULL},
     {"test_dict_api", test_dict_api, METH_NOARGS, _Py_NULL},
+    {"test_dict_pop", test_dict_pop, METH_NOARGS, _Py_NULL},
     {"test_long_api", test_long_api, METH_NOARGS, _Py_NULL},
 #ifdef TEST_MANAGED_DICT
     {"test_managed_dict", test_managed_dict, METH_NOARGS, _Py_NULL},
