gh-141510, PEP 814: Add frozendict support to pickle

Lib/test/picklecommon.py

@@ -243,6 +243,9 @@ class MyUnicode(str):
     class MyUnicode(unicode):
         sample = unicode(r"hello \u1234", "raw-unicode-escape")
 
+class DictKey:
+    pass
+
 class MyTuple(tuple):
     sample = (1, 2, 3)
 
@@ -263,6 +266,17 @@ class MyFrozenSet(frozenset):
              MyStr, MyUnicode,
              MyTuple, MyList, MyDict, MySet, MyFrozenSet]
 
+try:
+    frozendict
+except NameError:
+    # Python 3.14 and older
+    pass
+else:
+    class MyFrozenDict(dict):
+        sample = frozendict({"a": 1, "b": 2})
+    myclasses.append(MyFrozenDict)
+
+
 # For test_newobj_overridden_new
 class MyIntWithNew(int):
     def __new__(cls, value):

Lib/test/pickletester.py

@@ -2839,11 +2839,13 @@ def test_recursive_multi(self):
             self.assertEqual(list(x[0].attr.keys()), [1])
             self.assertIs(x[0].attr[1], x)
 
-    def _test_recursive_collection_and_inst(self, factory, oldminproto=None):
+    def _test_recursive_collection_and_inst(self, factory, oldminproto=None,
+                                            minprotocol=0):
         if self.py_version < (3, 0):
             self.skipTest('"classic" classes are not interoperable with Python 2')
         # Mutable object containing a collection containing the original
         # object.
+        protocols = range(minprotocol, pickle.HIGHEST_PROTOCOL + 1)
         o = Object()
         o.attr = factory([o])
         t = type(o.attr)
@@ -2883,6 +2885,11 @@ def test_recursive_tuple_and_inst(self):
     def test_recursive_dict_and_inst(self):
         self._test_recursive_collection_and_inst(dict.fromkeys, oldminproto=0)
 
+    def test_recursive_frozendict_and_inst(self):
+        if self.py_version < (3, 15):
+            self.skipTest('need frozendict')
+        self._test_recursive_collection_and_inst(frozendict.fromkeys, minprotocol=2)
+
     def test_recursive_set_and_inst(self):
         self._test_recursive_collection_and_inst(set)
 
@@ -2904,6 +2911,54 @@ def test_recursive_set_subclass_and_inst(self):
     def test_recursive_frozenset_subclass_and_inst(self):
         self._test_recursive_collection_and_inst(MyFrozenSet)
 
+    def _test_recursive_collection_in_key(self, factory, minprotocol=0):
+        protocols = range(minprotocol, pickle.HIGHEST_PROTOCOL + 1)
+        key = DictKey()
+        o = factory({key: 1})
+        key.attr = o
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                s = self.dumps(o, proto)
+                x = self.loads(s)
+                keys = list(x.keys())
+                self.assertEqual(len(keys), 1)
+                self.assertIs(keys[0].attr, x)
+
+    def test_recursive_dict_in_key(self):
+        self._test_recursive_collection_in_key(dict)
+
+    def test_recursive_dict_subclass_in_key(self):
+        self._test_recursive_collection_in_key(MyDict)
+
+    def test_recursive_frozendict_in_key(self):
+        self._test_recursive_collection_in_key(frozendict, minprotocol=2)
+
+    def test_recursive_frozendict_subclass_in_key(self):
+        self._test_recursive_collection_in_key(MyFrozenDict)
+
+    def _test_recursive_collection_in_value(self, factory, minprotocol=0):
+        protocols = range(minprotocol, pickle.HIGHEST_PROTOCOL + 1)
+        o = factory(key=[])
+        o['key'].append(o)
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                s = self.dumps(o, proto)
+                x = self.loads(s)
+                self.assertEqual(len(x['key']), 1)
+                self.assertIs(x['key'][0], x)
+
+    def test_recursive_dict_in_value(self):
+        self._test_recursive_collection_in_value(dict)
+
+    def test_recursive_dict_subclass_in_value(self):
+        self._test_recursive_collection_in_value(MyDict)
+
+    def test_recursive_frozendict_in_value(self):
+        self._test_recursive_collection_in_value(frozendict, minprotocol=2)
+
+    def test_recursive_frozendict_subclass_in_value(self):
+        self._test_recursive_collection_in_value(MyFrozenDict)
+
     def test_recursive_inst_state(self):
         # Mutable object containing itself.
         y = REX_state()

Lib/test/test_dict.py

@@ -1731,6 +1731,10 @@ class FrozenDict(frozendict):
     pass
 
 
+class FrozenDictSlots(frozendict):
+    __slots__ = ('slot_attr',)
+
+
 class FrozenDictTests(unittest.TestCase):
     def test_copy(self):
         d = frozendict(x=1, y=2)
@@ -1773,10 +1777,8 @@ def test_repr(self):
         d = frozendict(x=1, y=2)
         self.assertEqual(repr(d), "frozendict({'x': 1, 'y': 2})")
 
-        class MyFrozenDict(frozendict):
-            pass
-        d = MyFrozenDict(x=1, y=2)
-        self.assertEqual(repr(d), "MyFrozenDict({'x': 1, 'y': 2})")
+        d = FrozenDict(x=1, y=2)
+        self.assertEqual(repr(d), "FrozenDict({'x': 1, 'y': 2})")
 
     def test_hash(self):
         # hash() doesn't rely on the items order
@@ -1825,6 +1827,40 @@ def __new__(self):
         self.assertEqual(type(fd), DictSubclass)
         self.assertEqual(created, frozendict(x=1))
 
+    def test_pickle(self):
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            for fd in (
+                frozendict(),
+                frozendict(x=1, y=2),
+                FrozenDict(x=1, y=2),
+                FrozenDictSlots(x=1, y=2),
+            ):
+                if type(fd) == FrozenDict:
+                    fd.attr = 123
+                if type(fd) == FrozenDictSlots:
+                    fd.slot_attr = 456
+                with self.subTest(fd=fd, proto=proto):
+                    if proto >= 2:
+                        p = pickle.dumps(fd, proto)
+                        fd2 = pickle.loads(p)
+                        self.assertEqual(fd2, fd)
+                        self.assertEqual(type(fd2), type(fd))
+                        if type(fd) == FrozenDict:
+                            self.assertEqual(fd2.attr, 123)
+                        if type(fd) == FrozenDictSlots:
+                            self.assertEqual(fd2.slot_attr, 456)
+                    else:
+                        # protocol 0 and 1 don't support frozendict
+                        with self.assertRaises(TypeError):
+                            pickle.dumps(fd, proto)
+
+    def test_pickle_iter(self):
+        it = iter(frozendict(x=1, y=2))
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            p = pickle.dumps(it, proto)
+            it2 = pickle.loads(p)
+            self.assertEqual(list(it2), ['x', 'y'])
+
 
 if __name__ == "__main__":
     unittest.main()

Objects/dictobject.c

@@ -7930,6 +7930,18 @@ _PyObject_InlineValuesConsistencyCheck(PyObject *obj)
 
 // --- frozendict implementation ---------------------------------------------
 
+static PyObject *
+frozendict_getnewargs(PyObject *op, PyObject *Py_UNUSED(dummy))
+{
+    // Call dict(op): convert 'op' frozendict to a dict
+    PyObject *arg = PyObject_CallOneArg((PyObject*)&PyDict_Type, op);
+    if (arg == NULL) {
+        return NULL;
+    }
+    return Py_BuildValue("(N)", arg);
+}
+
+
 static PyNumberMethods frozendict_as_number = {
     .nb_or = frozendict_or,
 };
@@ -7951,6 +7963,7 @@ static PyMethodDef frozendict_methods[] = {
     DICT_COPY_METHODDEF
     DICT___REVERSED___METHODDEF
     {"__class_getitem__", Py_GenericAlias, METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
+    {"__getnewargs__", frozendict_getnewargs, METH_NOARGS},
     {NULL,              NULL}   /* sentinel */
 };