gh-111262: Add PyDict_Pop() function

Change the API of the internal _PyDict_Pop_KnownHash() function to return an int. Co-Authored-By: Stefan Behnel <stefan_ml@behnel.de>
python · vstinner · Nov 10, 2023 · Nov 12, 2023 · Nov 12, 2023 · Nov 12, 2023
commit b3e7de44f2ff653cf833bd45653ebda2a6630a3c
diff --git a/Doc/c-api/dict.rst b/Doc/c-api/dict.rst
@@ -173,6 +173,23 @@ Dictionary Objects
 
    .. versionadded:: 3.4
 
+
+.. c:function:: int PyDict_Pop(PyObject *p, PyObject *key, PyObject *default_value, PyObject *result)
+
+   Remove *key* from dictionary *p* and return the removed value.
+
+   - If the key is present, set *\*result* to a :term:`strong reference` to the
+     removed value, and return ``1``.
+   - If the key is missing and *default_value* is not NULL, set *\*result*
+     to a :term:`strong reference` to *default_value*, and return ``0``.
+   - If the key is missing and *default_value* is NULL, raise :exc:`KeyError`,
+     and return -1.
+   - On error, raise an exception and return ``-1``.
+
+   This is the same as :meth:`dict.pop`.
+
+   .. versionadded:: 3.13
+
 .. c:function:: PyObject* PyDict_Items(PyObject *p)
 
    Return a :c:type:`PyListObject` containing all the items from the dictionary.

diff --git a/Doc/data/stable_abi.dat b/Doc/data/stable_abi.dat
@@ -1164,6 +1164,10 @@ New Features
   :c:func:`PyErr_WriteUnraisable`, but allow to customize the warning mesage.
   (Contributed by Serhiy Storchaka in :gh:`108082`.)
 
+* Add :c:func:`PyDict_Pop` function: remove a key from a dictionary and return
+  the removed value. This is the same as :meth:`dict.pop`.
+  (Contributed by Victor Stinner in :gh:`111262`.)
+
 
 Porting to Python 3.13
 ----------------------

diff --git a/Include/dictobject.h b/Include/dictobject.h
@@ -66,6 +66,11 @@ PyAPI_FUNC(int) PyDict_DelItemString(PyObject *dp, const char *key);
 // - On error, raise an exception and return -1.
 PyAPI_FUNC(int) PyDict_GetItemRef(PyObject *mp, PyObject *key, PyObject **result);
 PyAPI_FUNC(int) PyDict_GetItemStringRef(PyObject *mp, const char *key, PyObject **result);
+
+PyAPI_FUNC(int) PyDict_Pop(PyObject *dict,
+    PyObject *key,
+    PyObject *default_value,
+    PyObject **result);
 #endif
 
 #if !defined(Py_LIMITED_API) || Py_LIMITED_API+0 >= 0x030A0000

diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
@@ -116,7 +116,12 @@ extern PyObject *_PyDict_LoadGlobal(PyDictObject *, PyDictObject *, PyObject *);
 extern int _PyDict_SetItem_Take2(PyDictObject *op, PyObject *key, PyObject *value);
 extern int _PyObjectDict_SetItem(PyTypeObject *tp, PyObject **dictptr, PyObject *name, PyObject *value);
 
-extern PyObject *_PyDict_Pop_KnownHash(PyObject *, PyObject *, Py_hash_t, PyObject *);
+extern int _PyDict_Pop_KnownHash(
+    PyDictObject *dict,
+    PyObject *key,
+    Py_hash_t hash,
+    PyObject *default_value,
+    PyObject **result);
 
 #define DKIX_EMPTY (-1)
 #define DKIX_DUMMY (-2)  /* Used internally */

diff --git a/Lib/test/test_capi/test_dict.py b/Lib/test/test_capi/test_dict.py
@@ -432,6 +432,49 @@ def test_dict_mergefromseq2(self):
         # CRASHES mergefromseq2({}, NULL, 0)
         # CRASHES mergefromseq2(NULL, {}, 0)
 
+    def test_dict_pop(self):
+        # Test PyDict_Pop()
+        dict_pop = _testcapi.dict_pop
+        default = object()
+
+        def expect_value(mydict, key, expected_value):
+            self.assertEqual(dict_pop(mydict.copy(), key, default),
+                             (1, expected_value))
+            self.assertEqual(dict_pop(mydict.copy(), key, NULL),
+                             (1, expected_value))
+
+        def check_default(mydict, key):
+            result, value = dict_pop(mydict, key, default)
+            self.assertIs(value, default)
+            self.assertEqual(result, 0)
+
+        # key present
+        mydict = {"key": 2, "key2": 5}
+        expect_value(mydict, "key", 2)
+        expect_value(mydict, "key2", 5)
+
+        # key missing
+        check_default({}, "key")  # empty dict has a fast path
+        check_default({"a": 1}, "key")
+        self.assertRaises(KeyError, dict_pop, {}, "key", NULL)
+        self.assertRaises(KeyError, dict_pop, {"a": 1}, "key", NULL)
+
+        # dict error
+        not_dict = "string"
+        self.assertRaises(SystemError, dict_pop, not_dict, "key", default)
+
+        # key error
+        not_hashable_key = ["list"]
+        check_default({}, not_hashable_key)  # don't hash key if dict is empty
+        with self.assertRaises(TypeError):
+            dict_pop({'key': 1}, not_hashable_key, NULL)
+        with self.assertRaises(TypeError):
+            dict_pop({'key': 1}, not_hashable_key, default)
+        dict_pop({}, NULL, default)  # don't check key if dict is empty
+
+        # CRASHES dict_pop(NULL, "key")
+        # CRASHES dict_pop({"a": 1}, NULL, default)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/test/test_stable_abi_ctypes.py b/Lib/test/test_stable_abi_ctypes.py
diff --git a/Misc/NEWS.d/next/C API/2023-11-10-10-21-38.gh-issue-111262.2utB5m.rst b/Misc/NEWS.d/next/C API/2023-11-10-10-21-38.gh-issue-111262.2utB5m.rst
@@ -0,0 +1,3 @@
+Add :c:func:`PyDict_Pop` function: remove a key from a dictionary and return
+the removed value. This is the same as :meth:`dict.pop`. Patch by Stefan
+Behnel and Victor Stinner.
@@ -2481,3 +2481,5 @@
 [function._Py_SetRefcnt]
     added = '3.13'
     abi_only = true
+[function.PyDict_Pop]
+    added = '3.13'
@@ -1087,8 +1087,8 @@ bounded_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwds
        The cache dict holds one reference to the link.
        We created one other reference when the link was created.
        The linked list only has borrowed references. */
-    popresult = _PyDict_Pop_KnownHash(self->cache, link->key,
-                                      link->hash, Py_None);
+    (void)_PyDict_Pop_KnownHash((PyDictObject*)self->cache, link->key,
+                                link->hash, Py_None, &popresult);
     if (popresult == Py_None) {
         /* Getting here means that the user function call or another
            thread has already removed the old key from the dictionary.

diff --git a/Modules/_testcapi/dict.c b/Modules/_testcapi/dict.c
@@ -331,6 +331,25 @@ dict_mergefromseq2(PyObject *self, PyObject *args)
 }
 
 
+static PyObject *
+dict_pop(PyObject *self, PyObject *args)
+{
+    PyObject *dict, *key, *default_value;
+    if (!PyArg_ParseTuple(args, "OOO", &dict, &key, &default_value)) {
+        return NULL;
+    }
+    NULLABLE(dict);
+    NULLABLE(key);
+    NULLABLE(default_value);
+    PyObject *result = UNINITIALIZED_PTR;
+    int res = PyDict_Pop(dict, key, default_value, &result);
+    if (result == NULL) {
+        return NULL;
+    }
+    return Py_BuildValue("iN", res, result);
+}
+
+
 static PyMethodDef test_methods[] = {
     {"dict_check", dict_check, METH_O},
     {"dict_checkexact", dict_checkexact, METH_O},
@@ -358,7 +377,7 @@ static PyMethodDef test_methods[] = {
     {"dict_merge", dict_merge, METH_VARARGS},
     {"dict_update", dict_update, METH_VARARGS},
     {"dict_mergefromseq2", dict_mergefromseq2, METH_VARARGS},
-
+    {"dict_pop", dict_pop, METH_VARARGS},
     {NULL},
 };
 

@@ -2226,64 +2226,94 @@ PyDict_Next(PyObject *op, Py_ssize_t *ppos, PyObject **pkey, PyObject **pvalue)
     return _PyDict_Next(op, ppos, pkey, pvalue, NULL);
 }
 
+
 /* Internal version of dict.pop(). */
-PyObject *
-_PyDict_Pop_KnownHash(PyObject *dict, PyObject *key, Py_hash_t hash, PyObject *deflt)
+int
+_PyDict_Pop_KnownHash(PyDictObject *mp, PyObject *key, Py_hash_t hash,
+                      PyObject *default_value, PyObject **result)
 {
-    Py_ssize_t ix;
-    PyObject *old_value;
-    PyDictObject *mp;
-    PyInterpreterState *interp = _PyInterpreterState_GET();
-
-    assert(PyDict_Check(dict));
-    mp = (PyDictObject *)dict;
+    assert(PyDict_Check(mp));
 
     if (mp->ma_used == 0) {
-        if (deflt) {
-            return Py_NewRef(deflt);
+        if (default_value) {
+            *result = Py_NewRef(default_value);
+            return 0;
         }
+        *result = NULL;
         _PyErr_SetKeyError(key);
-        return NULL;
+        return -1;
     }
-    ix = _Py_dict_lookup(mp, key, hash, &old_value);
-    if (ix == DKIX_ERROR)
-        return NULL;
+
+    PyObject *old_value;
+    Py_ssize_t ix = _Py_dict_lookup(mp, key, hash, &old_value);
+    if (ix == DKIX_ERROR) {
+        *result = NULL;
+        return -1;
+    }
+
     if (ix == DKIX_EMPTY || old_value == NULL) {
-        if (deflt) {
-            return Py_NewRef(deflt);
+        if (default_value) {
+            *result = Py_NewRef(default_value);
+            return 0;
         }
+        *result = NULL;
         _PyErr_SetKeyError(key);
-        return NULL;
+        return -1;
     }
+
     assert(old_value != NULL);
+    PyInterpreterState *interp = _PyInterpreterState_GET();
     uint64_t new_version = _PyDict_NotifyEvent(
             interp, PyDict_EVENT_DELETED, mp, key, NULL);
     delitem_common(mp, hash, ix, Py_NewRef(old_value), new_version);
 
     ASSERT_CONSISTENT(mp);
-    return old_value;
+    *result = old_value;
+    return 1;
 }
 
-PyObject *
-_PyDict_Pop(PyObject *dict, PyObject *key, PyObject *deflt)
+
+int
+PyDict_Pop(PyObject *op, PyObject *key, PyObject *default_value, PyObject **result)
 {
-    Py_hash_t hash;
+    if (!PyDict_Check(op)) {
+        *result = NULL;
+        PyErr_BadInternalCall();
+        return -1;
+    }
+    PyDictObject *dict = (PyDictObject *)op;
 
-    if (((PyDictObject *)dict)->ma_used == 0) {
-        if (deflt) {
-            return Py_NewRef(deflt);
+    if (dict->ma_used == 0) {
+        if (default_value) {
+            *result = Py_NewRef(default_value);
+            return 0;
         }
+        *result = NULL;
         _PyErr_SetKeyError(key);
-        return NULL;
+        return -1;
     }
+
+    Py_hash_t hash;
     if (!PyUnicode_CheckExact(key) || (hash = unicode_get_hash(key)) == -1) {
         hash = PyObject_Hash(key);
-        if (hash == -1)
-            return NULL;
+        if (hash == -1) {
+            *result = NULL;
+            return -1;
+        }
     }
-    return _PyDict_Pop_KnownHash(dict, key, hash, deflt);
+    return _PyDict_Pop_KnownHash(dict, key, hash, default_value, result);
 }
 
+
+PyObject *
+_PyDict_Pop(PyObject *dict, PyObject *key, PyObject *default_value)
+{
+    PyObject *result;
+    (void)PyDict_Pop(dict, key, default_value, &result);
+    return result;
+}
+
+
 /* Internal version of dict.from_keys().  It is subclass-friendly. */
 PyObject *
 _PyDict_FromKeys(PyObject *cls, PyObject *iterable, PyObject *value)

diff --git a/Objects/odictobject.c b/Objects/odictobject.c
@@ -1049,7 +1049,7 @@ _odict_popkey_hash(PyObject *od, PyObject *key, PyObject *failobj,
             return NULL;
         }
         /* Now delete the value from the dict. */
-        value = _PyDict_Pop_KnownHash(od, key, hash, failobj);
+        (void)_PyDict_Pop_KnownHash((PyDictObject *)od, key, hash, failobj, &value);
     }
     else if (value == NULL && !PyErr_Occurred()) {
         /* Apply the fallback value, if necessary. */