From fa618e6aae9bb63dccc383e5e4026bc596f7d388 Mon Sep 17 00:00:00 2001
From: Victor Stinner <vstinner@python.org>
Date: Mon, 13 Nov 2023 14:04:06 +0100
Subject: [PATCH] Add tests

---
 Lib/test/test_capi/test_dict.py | 43 +++++++++++++++++++++++++++++---
 Modules/_testcapi/dict.c        | 44 +++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/Lib/test/test_capi/test_dict.py b/Lib/test/test_capi/test_dict.py
index 4182dbad249753c..e5a36392162637c 100644
--- a/Lib/test/test_capi/test_dict.py
+++ b/Lib/test/test_capi/test_dict.py
@@ -435,21 +435,34 @@ def test_dict_mergefromseq2(self):
     def test_dict_pop(self):
         # Test PyDict_Pop()
         dict_pop = _testcapi.dict_pop
+        dict_pop_null = _testcapi.dict_pop_null
 
-        # key present
+        # key present, get removed value
         mydict = {"key": "value", "key2": "value2"}
         self.assertEqual(dict_pop(mydict, "key"), (1, "value"))
         self.assertEqual(mydict, {"key2": "value2"})
         self.assertEqual(dict_pop(mydict, "key2"), (1, "value2"))
         self.assertEqual(mydict, {})
 
-        # key missing; empty dict has a fast path
+        # key present, ignore removed value
+        mydict = {"key": "value", "key2": "value2"}
+        self.assertEqual(dict_pop_null(mydict, "key"), 1)
+        self.assertEqual(mydict, {"key2": "value2"})
+        self.assertEqual(dict_pop_null(mydict, "key2"), 1)
+        self.assertEqual(mydict, {})
+
+        # key missing, expect removed value; empty dict has a fast path
         self.assertEqual(dict_pop({}, "key"), (0, NULL))
         self.assertEqual(dict_pop({"a": 1}, "key"), (0, NULL))
 
+        # key missing, ignored removed value; empty dict has a fast path
+        self.assertEqual(dict_pop_null({}, "key"), 0)
+        self.assertEqual(dict_pop_null({"a": 1}, "key"), 0)
+
         # dict error
-        not_dict = "string"
+        not_dict = UserDict({1: 2})
         self.assertRaises(SystemError, dict_pop, not_dict, "key")
+        self.assertRaises(SystemError, dict_pop_null, not_dict, "key")
 
         # key error; don't hash key if dict is empty
         not_hashable_key = ["list"]
@@ -459,7 +472,29 @@ def test_dict_pop(self):
         dict_pop({}, NULL)  # key is not checked if dict is empty
 
         # CRASHES dict_pop(NULL, "key")
-        # CRASHES dict_pop({"a": 1}, NULL, default)
+        # CRASHES dict_pop({"a": 1}, NULL)
+
+    def test_dict_popstring(self):
+        # Test PyDict_PopString()
+        dict_popstring = _testcapi.dict_popstring
+
+        # key present
+        mydict = {"key": "value", "key2": "value2"}
+        self.assertEqual(dict_popstring(mydict, "key"), (1, "value"))
+        self.assertEqual(mydict, {"key2": "value2"})
+        self.assertEqual(dict_popstring(mydict, "key2"), (1, "value2"))
+        self.assertEqual(mydict, {})
+
+        # key missing; empty dict has a fast path
+        self.assertEqual(dict_popstring({}, "key"), (0, NULL))
+        self.assertEqual(dict_popstring({"a": 1}, "key"), (0, NULL))
+
+        # dict error
+        not_dict = UserDict({1: 2})
+        self.assertRaises(SystemError, dict_popstring, not_dict, "key")
+
+        # CRASHES dict_popstring(NULL, "key")
+        # CRASHES dict_popstring({"a": 1}, NULL)
 
 
 if __name__ == "__main__":
diff --git a/Modules/_testcapi/dict.c b/Modules/_testcapi/dict.c
index 3590a9ee8bdaacb..26d86cfa8c8954f 100644
--- a/Modules/_testcapi/dict.c
+++ b/Modules/_testcapi/dict.c
@@ -334,6 +334,7 @@ dict_mergefromseq2(PyObject *self, PyObject *args)
 static PyObject *
 dict_pop(PyObject *self, PyObject *args)
 {
+    // Test PyDict_Pop(dict, key, &value)
     PyObject *dict, *key;
     if (!PyArg_ParseTuple(args, "OO", &dict, &key)) {
         return NULL;
@@ -353,6 +354,47 @@ dict_pop(PyObject *self, PyObject *args)
 }
 
 
+static PyObject *
+dict_pop_null(PyObject *self, PyObject *args)
+{
+    // Test PyDict_Pop(dict, key, NULL)
+    PyObject *dict, *key;
+    if (!PyArg_ParseTuple(args, "OO", &dict, &key)) {
+        return NULL;
+    }
+    NULLABLE(dict);
+    NULLABLE(key);
+    int res = PyDict_Pop(dict, key,  NULL);
+    if (res < 0) {
+        return NULL;
+    }
+    return PyLong_FromLong(res);
+}
+
+
+static PyObject *
+dict_popstring(PyObject *self, PyObject *args)
+{
+    PyObject *dict;
+    const char *key;
+    Py_ssize_t key_size;
+    if (!PyArg_ParseTuple(args, "Oz#", &dict, &key, &key_size)) {
+        return NULL;
+    }
+    NULLABLE(dict);
+    PyObject *result = UNINITIALIZED_PTR;
+    int res = PyDict_PopString(dict, key,  &result);
+    if (res < 0) {
+        assert(result == NULL);
+        return NULL;
+    }
+    if (result == NULL) {
+        result = Py_NewRef(Py_None);
+    }
+    return Py_BuildValue("iN", res, result);
+}
+
+
 static PyMethodDef test_methods[] = {
     {"dict_check", dict_check, METH_O},
     {"dict_checkexact", dict_checkexact, METH_O},
@@ -381,6 +423,8 @@ static PyMethodDef test_methods[] = {
     {"dict_update", dict_update, METH_VARARGS},
     {"dict_mergefromseq2", dict_mergefromseq2, METH_VARARGS},
     {"dict_pop", dict_pop, METH_VARARGS},
+    {"dict_pop_null", dict_pop_null, METH_VARARGS},
+    {"dict_popstring", dict_popstring, METH_VARARGS},
     {NULL},
 };