gh-141510: Return frozendict unmodified in PyDict_Copy() (#145505)

Add also the internal _PyDict_CopyAsDict() function.

Author: vstinner

Include/internal/pycore_dict.h

@@ -160,6 +160,8 @@ extern void _PyDict_Clear_LockHeld(PyObject *op);
 PyAPI_FUNC(void) _PyDict_EnsureSharedOnRead(PyDictObject *mp);
 #endif
 
+extern PyObject* _PyDict_CopyAsDict(PyObject *op);
+
 #define DKIX_EMPTY (-1)
 #define DKIX_DUMMY (-2)  /* Used internally */
 #define DKIX_ERROR (-3)

Lib/test/test_capi/test_dict.py

@@ -98,14 +98,18 @@ def test_dict_copy(self):
         # Test PyDict_Copy()
         copy = _testlimitedcapi.dict_copy
         for dict_type in ANYDICT_TYPES:
-            if issubclass(dict_type, frozendict):
-                expected_type = frozendict
-            else:
-                expected_type = dict
             dct = dict_type({1: 2})
             dct_copy = copy(dct)
-            self.assertIs(type(dct_copy), expected_type)
-            self.assertEqual(dct_copy, dct)
+            if dict_type == frozendict:
+                expected_type = frozendict
+                self.assertIs(dct_copy, dct)
+            else:
+                if issubclass(dict_type, frozendict):
+                    expected_type = frozendict
+                else:
+                    expected_type = dict
+                self.assertIs(type(dct_copy), expected_type)
+                self.assertEqual(dct_copy, dct)
 
         for test_type in NOT_ANYDICT_TYPES + OTHER_TYPES:
             self.assertRaises(SystemError, copy, test_type())

Objects/dictobject.c

@@ -4235,9 +4235,6 @@ static PyObject *
 dict_copy_impl(PyDictObject *self)
 /*[clinic end generated code: output=ffb782cf970a5c39 input=73935f042b639de4]*/
 {
-    if (PyFrozenDict_CheckExact(self)) {
-        return Py_NewRef(self);
-    }
     return PyDict_Copy((PyObject *)self);
 }
 
@@ -4263,18 +4260,17 @@ copy_values(PyDictValues *values)
 }
 
 static PyObject *
-copy_lock_held(PyObject *o)
+copy_lock_held(PyObject *o, int as_frozendict)
 {
     PyObject *copy;
     PyDictObject *mp;
-    int frozendict = PyFrozenDict_Check(o);
 
     ASSERT_DICT_LOCKED(o);
 
     mp = (PyDictObject *)o;
     if (mp->ma_used == 0) {
         /* The dict is empty; just return a new dict. */
-        if (frozendict) {
+        if (as_frozendict) {
             return PyFrozenDict_New(NULL);
         }
         else {
@@ -4288,7 +4284,7 @@ copy_lock_held(PyObject *o)
         if (newvalues == NULL) {
             return PyErr_NoMemory();
         }
-        if (frozendict) {
+        if (as_frozendict) {
             split_copy = (PyDictObject *)PyObject_GC_New(PyFrozenDictObject,
                                                          &PyFrozenDict_Type);
         }
@@ -4307,7 +4303,7 @@ copy_lock_held(PyObject *o)
         split_copy->ma_used = mp->ma_used;
         split_copy->_ma_watcher_tag = 0;
         dictkeys_incref(mp->ma_keys);
-        if (frozendict) {
+        if (as_frozendict) {
             PyFrozenDictObject *frozen = (PyFrozenDictObject *)split_copy;
             frozen->ma_hash = -1;
         }
@@ -4318,7 +4314,7 @@ copy_lock_held(PyObject *o)
     if (Py_TYPE(mp)->tp_iter == dict_iter &&
             mp->ma_values == NULL &&
             (mp->ma_used >= (mp->ma_keys->dk_nentries * 2) / 3) &&
-            !frozendict)
+            !as_frozendict)
     {
         /* Use fast-copy if:
 
@@ -4350,7 +4346,7 @@ copy_lock_held(PyObject *o)
         return (PyObject *)new;
     }
 
-    if (frozendict) {
+    if (as_frozendict) {
         copy = PyFrozenDict_New(NULL);
     }
     else {
@@ -4364,6 +4360,19 @@ copy_lock_held(PyObject *o)
     return NULL;
 }
 
+// Similar to PyDict_Copy(), but copy also frozendict.
+static PyObject *
+_PyDict_Copy(PyObject *o)
+{
+    assert(PyAnyDict_Check(o));
+
+    PyObject *res;
+    Py_BEGIN_CRITICAL_SECTION(o);
+    res = copy_lock_held(o, PyFrozenDict_Check(o));
+    Py_END_CRITICAL_SECTION();
+    return res;
+}
+
 PyObject *
 PyDict_Copy(PyObject *o)
 {
@@ -4372,11 +4381,22 @@ PyDict_Copy(PyObject *o)
         return NULL;
     }
 
-    PyObject *res;
-    Py_BEGIN_CRITICAL_SECTION(o);
+    if (PyFrozenDict_CheckExact(o)) {
+        return Py_NewRef(o);
+    }
+
+    return _PyDict_Copy(o);
+}
 
-    res = copy_lock_held(o);
+// Similar to PyDict_Copy(), but return a dict if the argument is a frozendict.
+PyObject*
+_PyDict_CopyAsDict(PyObject *o)
+{
+    assert(PyAnyDict_Check(o));
 
+    PyObject *res;
+    Py_BEGIN_CRITICAL_SECTION(o);
+    res = copy_lock_held(o, 0);
     Py_END_CRITICAL_SECTION();
     return res;
 }
@@ -4925,7 +4945,7 @@ dict_or(PyObject *self, PyObject *other)
     if (!PyAnyDict_Check(self) || !PyAnyDict_Check(other)) {
         Py_RETURN_NOTIMPLEMENTED;
     }
-    PyObject *new = PyDict_Copy(self);
+    PyObject *new = _PyDict_Copy(self);
     if (new == NULL) {
         return NULL;
     }
@@ -6479,7 +6499,7 @@ dictitems_xor_lock_held(PyObject *d1, PyObject *d2)
     ASSERT_DICT_LOCKED(d1);
     ASSERT_DICT_LOCKED(d2);
 
-    PyObject *temp_dict = copy_lock_held(d1);
+    PyObject *temp_dict = copy_lock_held(d1, PyFrozenDict_Check(d1));
     if (temp_dict == NULL) {
         return NULL;
     }

Objects/typeobject.c

@@ -4850,21 +4850,9 @@ type_new_get_slots(type_new_ctx *ctx, PyObject *dict)
 static PyTypeObject*
 type_new_init(type_new_ctx *ctx)
 {
-    PyObject *dict;
-    if (PyFrozenDict_Check(ctx->orig_dict)) {
-        dict = PyDict_New();
-        if (dict == NULL) {
-            goto error;
-        }
-        if (PyDict_Merge(dict, ctx->orig_dict, 1) < 0) {
-            goto error;
-        }
-    }
-    else {
-        dict = PyDict_Copy(ctx->orig_dict);
-        if (dict == NULL) {
-            goto error;
-        }
+    PyObject *dict = _PyDict_CopyAsDict(ctx->orig_dict);
+    if (dict == NULL) {
+        goto error;
     }
 
     if (type_new_get_slots(ctx, dict) < 0) {