decouple error checking and kw merging + tests

dg-pb · Jan 6, 2025 · 106e6bc · 106e6bc
1 parent e8d7521
commit 106e6bc
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 19 deletions.
diff --git a/Lib/functools.py b/Lib/functools.py
@@ -326,8 +326,9 @@ def _partial_new(cls, func, /, *args, **keywords):
                             "or a descriptor")
     if args and args[-1] is Placeholder:
         raise TypeError("trailing Placeholders are not allowed")
-    if keywords and Placeholder in keywords.values():
-        raise TypeError("Placeholder cannot be passed as a keyword argument")
+    for value in keywords.values():
+        if value is Placeholder:
+            raise TypeError("Placeholder cannot be passed as a keyword argument")
     if isinstance(func, base_cls):
         pto_phcount = func._phcount
         tot_args = func.args

diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py
@@ -233,6 +233,12 @@ def test_placeholders(self):
         actual_args, actual_kwds = p('x', 'y')
         self.assertEqual(actual_args, ('x', 0, 'y', 1))
         self.assertEqual(actual_kwds, {})
+        # Checks via `is` and not `eq`
+        # thus unittest.mock.ANY isn't treated as Placeholder
+        p = self.partial(capture, unittest.mock.ANY)
+        actual_args, actual_kwds = p()
+        self.assertEqual(actual_args, (unittest.mock.ANY,))
+        self.assertEqual(actual_kwds, {})
 
     def test_placeholders_optimization(self):
         PH = self.module.Placeholder
@@ -253,6 +259,11 @@ def test_placeholders_kw_restriction(self):
         PH = self.module.Placeholder
         with self.assertRaisesRegex(TypeError, "Placeholder"):
             self.partial(capture, a=PH)
+        # Passes, as checks via `is` and not `eq`
+        p = self.partial(capture, a=unittest.mock.ANY)
+        actual_args, actual_kwds = p()
+        self.assertEqual(actual_args, ())
+        self.assertEqual(actual_kwds, {'a': unittest.mock.ANY})
 
     def test_construct_placeholder_singleton(self):
         PH = self.module.Placeholder

diff --git a/Modules/_functoolsmodule.c b/Modules/_functoolsmodule.c
@@ -196,6 +196,19 @@ partial_new(PyTypeObject *type, PyObject *args, PyObject *kw)
         return NULL;
     }
 
+    /* keyword Placeholder prohibition */
+    if (kw != NULL) {
+        PyObject *key, *val;
+        Py_ssize_t pos = 0;
+        while (PyDict_Next(kw, &pos, &key, &val)) {
+            if (val == phold) {
+                PyErr_SetString(PyExc_TypeError,
+                                "Placeholder cannot be passed as a keyword argument");
+                return NULL;
+            }
+        }
+    }
+
     /* check wrapped function / object */
     pto_args = pto_kw = NULL;
     int res = PyObject_TypeCheck(func, state->partial_type);
@@ -281,31 +294,29 @@ partial_new(PyTypeObject *type, PyObject *args, PyObject *kw)
     }
 
     if (pto_kw == NULL || PyDict_GET_SIZE(pto_kw) == 0) {
-        pto->kw = PyDict_New();
+        if (kw == NULL) {
+            pto->kw = PyDict_New();
+        }
+        else if (Py_REFCNT(kw) == 1) {
+            pto->kw = Py_NewRef(kw);
+        }
+        else {
+            pto->kw = PyDict_Copy(kw);
+        }
     }
     else {
         pto->kw = PyDict_Copy(pto_kw);
-    }
-    if (pto->kw == NULL) {
-        Py_DECREF(pto);
-        return NULL;
-    }
-    if (kw != NULL) {
-        PyObject *key, *val;
-        Py_ssize_t pos = 0;
-        while (PyDict_Next(kw, &pos, &key, &val)) {
-            if (val == phold) {
-                Py_DECREF(pto);
-                PyErr_SetString(PyExc_TypeError,
-                                "Placeholder cannot be passed as a keyword argument");
-                return NULL;
-            }
-            if (PyDict_SetItem(pto->kw, key, val) < 0) {
+        if (kw != NULL && pto->kw != NULL) {
+            if (PyDict_Merge(pto->kw, kw, 1) != 0) {
                 Py_DECREF(pto);
                 return NULL;
             }
         }
     }
+    if (pto->kw == NULL) {
+        Py_DECREF(pto);
+        return NULL;
+    }
 
     partial_setvectorcall(pto);
     return (PyObject *)pto;