gh-141510: Fix frozendict.fromkeys() for subclasses

Copy the frozendict.

Author: vstinner

Lib/test/test_dict.py

@@ -1787,6 +1787,31 @@ def test_hash(self):
         with self.assertRaisesRegex(TypeError, "unhashable type: 'list'"):
             hash(fd)
 
+    def test_fromkeys(self):
+        self.assertEqual(frozendict.fromkeys('abc'),
+                         frozendict(a=None, b=None, c=None))
+
+        # Subclass which overrides the constructor
+        class FrozenDictSubclass(frozendict):
+            def __new__(self):
+                return frozendict(x=1)
+
+        fd = FrozenDictSubclass.fromkeys("abc")
+        self.assertEqual(fd, frozendict(x=1, a=None, b=None, c=None))
+        self.assertEqual(type(fd), FrozenDictSubclass)
+
+        fd = FrozenDictSubclass.fromkeys(frozendict(y=2))
+        self.assertEqual(fd, frozendict(x=1, y=None))
+        self.assertEqual(type(fd), FrozenDictSubclass)
+
+        # Subclass which doesn't override the constructor
+        class FrozenDictSubclass2(frozendict):
+            pass
+
+        fd = FrozenDictSubclass2.fromkeys("abc")
+        self.assertEqual(fd, frozendict(a=None, b=None, c=None))
+        self.assertEqual(type(fd), FrozenDictSubclass2)
+
 
 if __name__ == "__main__":
     unittest.main()

Objects/dictobject.c

@@ -138,6 +138,7 @@ As a consequence of this, split keys have a maximum size of 16.
 // Forward declarations
 static PyObject* frozendict_new(PyTypeObject *type, PyObject *args,
                                 PyObject *kwds);
+static int dict_merge(PyObject *a, PyObject *b, int override);
 
 
 /*[clinic input]
@@ -3286,9 +3287,31 @@ _PyDict_FromKeys(PyObject *cls, PyObject *iterable, PyObject *value)
     int status;
 
     d = _PyObject_CallNoArgs(cls);
-    if (d == NULL)
+    if (d == NULL) {
         return NULL;
+    }
 
+    // If cls is a frozendict subclass with overridden constructor,
+    // copy the frozendict.
+    PyTypeObject *cls_type = _PyType_CAST(cls);
+    if (PyFrozenDict_Check(d)
+        && PyObject_IsSubclass(cls, (PyObject*)&PyFrozenDict_Type)
+        && cls_type->tp_new != frozendict_new)
+    {
+        // Subclass-friendly copy
+        PyObject *copy = frozendict_new(cls_type, NULL, NULL);
+        if (copy == NULL) {
+            Py_DECREF(d);
+            return NULL;
+        }
+        if (dict_merge(copy, d, 1) < 0) {
+            Py_DECREF(d);
+            Py_DECREF(copy);
+            return NULL;
+        }
+        Py_SETREF(d, copy);
+    }
+    assert(!PyFrozenDict_Check(d) || Py_REFCNT(d) == 1);
 
     if (PyDict_CheckExact(d)) {
         if (PyDict_CheckExact(iterable)) {
@@ -3359,7 +3382,7 @@ _PyDict_FromKeys(PyObject *cls, PyObject *iterable, PyObject *value)
 dict_iter_exit:;
         Py_END_CRITICAL_SECTION();
     }
-    else if (PyFrozenDict_CheckExact(d)) {
+    else if (PyFrozenDict_Check(d)) {
         while ((key = PyIter_Next(it)) != NULL) {
             // anydict_setitem_take2 consumes a reference to key
             status = anydict_setitem_take2((PyDictObject *)d,
@@ -7994,6 +8017,8 @@ frozendict_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     if (d == NULL) {
         return NULL;
     }
+    assert(Py_REFCNT(d) == 1);
+
     PyFrozenDictObject *self = _PyFrozenDictObject_CAST(d);
     self->ma_hash = -1;