Add more tests

vstinner · vstinner · commit 7c84afbf30ca · 2023-11-14T00:03:34.000+01:00
diff --git a/Lib/test/test_capi/test_dict.py b/Lib/test/test_capi/test_dict.py
@@ -477,23 +477,46 @@ def test_dict_pop(self):
     def test_dict_popstring(self):
         # Test PyDict_PopString()
         dict_popstring = _testcapi.dict_popstring
+        dict_popstring_null = _testcapi.dict_popstring_null
 
-        # key present
+        # key present, get removed value
         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 present, ignore removed value
+        mydict = {"key": "value", "key2": "value2"}
+        self.assertEqual(dict_popstring_null(mydict, "key"), 1)
+        self.assertEqual(mydict, {"key2": "value2"})
+        self.assertEqual(dict_popstring_null(mydict, "key2"), 1)
+        self.assertEqual(mydict, {})
+
         # key missing; empty dict has a fast path
         self.assertEqual(dict_popstring({}, "key"), (0, NULL))
+        self.assertEqual(dict_popstring_null({}, "key"), 0)
         self.assertEqual(dict_popstring({"a": 1}, "key"), (0, NULL))
+        self.assertEqual(dict_popstring_null({"a": 1}, "key"), 0)
+
+        # non-ASCII key
+        non_ascii = '\U0001f40d'
+        dct = {'\U0001f40d': 123}
+        self.assertEqual(dict_popstring(dct, '\U0001f40d'.encode()), (1, 123))
+        dct = {'\U0001f40d': 123}
+        self.assertEqual(dict_popstring_null(dct, '\U0001f40d'.encode()), 1)
 
         # dict error
         not_dict = UserDict({1: 2})
         self.assertRaises(SystemError, dict_popstring, not_dict, "key")
+        self.assertRaises(SystemError, dict_popstring_null, not_dict, "key")
+
+        # key error
+        self.assertRaises(UnicodeDecodeError, dict_popstring, {1: 2}, INVALID_UTF8)
+        self.assertRaises(UnicodeDecodeError, dict_popstring_null, {1: 2}, INVALID_UTF8)
 
         # CRASHES dict_popstring(NULL, "key")
+        # CRASHES dict_popstring({}, NULL)
         # CRASHES dict_popstring({"a": 1}, NULL)
 
 
diff --git a/Modules/_testcapi/dict.c b/Modules/_testcapi/dict.c
@@ -368,11 +368,7 @@ dict_pop_null(PyObject *self, PyObject *args)
     }
     NULLABLE(dict);
     NULLABLE(key);
-    int res = PyDict_Pop(dict, key,  NULL);
-    if (res < 0) {
-        return NULL;
-    }
-    return PyLong_FromLong(res);
+    RETURN_INT(PyDict_Pop(dict, key,  NULL));
 }
 
 
@@ -403,6 +399,20 @@ dict_popstring(PyObject *self, PyObject *args)
 }
 
 
+static PyObject *
+dict_popstring_null(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);
+    RETURN_INT(PyDict_PopString(dict, key,  NULL));
+}
+
+
 static PyMethodDef test_methods[] = {
     {"dict_check", dict_check, METH_O},
     {"dict_checkexact", dict_checkexact, METH_O},
@@ -433,6 +443,7 @@ static PyMethodDef test_methods[] = {
     {"dict_pop", dict_pop, METH_VARARGS},
     {"dict_pop_null", dict_pop_null, METH_VARARGS},
     {"dict_popstring", dict_popstring, METH_VARARGS},
+    {"dict_popstring_null", dict_popstring_null, METH_VARARGS},
     {NULL},
 };
 
