Add frozendict support to _PyCode_ConstantKey

Author: lgeiger

Lib/test/test_compile.py

@@ -776,6 +776,13 @@ def check_constant(self, func, expected):
             self.fail("unable to find constant %r in %r"
                       % (expected, func.__code__.co_consts))
 
+    @staticmethod
+    def _frozen_dict_consts(*consts):
+        """Use AST to make frozendict constants since it has no literal syntax"""
+        m = ast.Interactive([ast.Expr(ast.Constant(c)) for c in consts])
+        ast.fix_missing_locations(m)
+        return compile(m, "<test>", "single")
+
     # Merging equal constants is not a strict requirement for the Python
     # semantics, it's a more an implementation detail.
     @support.cpython_only
@@ -820,6 +827,21 @@ def check_same_constant(const):
         self.check_constant(f1, frozenset({0}))
         self.assertTrue(f1(0))
 
+        # two identical frozendicts merge into one constant
+        c = self._frozen_dict_consts(frozendict({0: 1}), frozendict({0: 1}))
+        self.assertEqual(c.co_consts, (frozendict({0: 1}),))
+
+        # empty frozendicts also merge
+        c = self._frozen_dict_consts(frozendict(), frozendict())
+        self.assertEqual(c.co_consts, (frozendict(),))
+
+        # frozendicts containing a nested frozendict value merge
+        c = self._frozen_dict_consts(
+            frozendict({0: frozendict({1: 2})}),
+            frozendict({0: frozendict({1: 2})}),
+        )
+        self.assertEqual(c.co_consts, (frozendict({0: frozendict({1: 2})}),))
+
     # Merging equal co_linetable is not a strict requirement
     # for the Python semantics, it's a more an implementation detail.
     @support.cpython_only
@@ -1033,6 +1055,16 @@ def check_different_constants(const1, const2):
         self.assertTrue(f1(0))
         self.assertTrue(f2(0.0))
 
+        # frozendicts with type-distinct keys must not merge (0 vs 0.0)
+        c = self._frozen_dict_consts(frozendict({0: 1}), frozendict({0.0: 1}))
+        self.assertEqual(c.co_consts, (frozendict({0: 1}), frozendict({0.0: 1})))
+        self.assertIsNot(c.co_consts[0], c.co_consts[1])
+
+        # frozendicts with type-distinct values must not merge (1 vs 1.0)
+        c = self._frozen_dict_consts(frozendict({0: 1}), frozendict({0: 1.0}))
+        self.assertEqual(c.co_consts, (frozendict({0: 1}), frozendict({0: 1.0})))
+        self.assertIsNot(c.co_consts[0], c.co_consts[1])
+
     def test_path_like_objects(self):
         # An implicit test for PyUnicode_FSDecoder().
         compile("42", FakePath("test_compile_pathlike"), "single")

Objects/codeobject.c

@@ -3039,6 +3039,45 @@ _PyCode_ConstantKey(PyObject *op)
         Py_DECREF(set);
         return key;
     }
+    else if (PyFrozenDict_CheckExact(op)) {
+        PyObject *dict = PyDict_New();
+        if (dict == NULL) {
+            return NULL;
+        }
+
+        Py_ssize_t pos = 0;
+        PyObject *k_obj, *v_obj;
+        while (PyDict_Next(op, &pos, &k_obj, &v_obj)) {
+            PyObject *k_key = _PyCode_ConstantKey(k_obj);
+            if (k_key == NULL) {
+                Py_DECREF(dict);
+                return NULL;
+            }
+            PyObject *v_key = _PyCode_ConstantKey(v_obj);
+            if (v_key == NULL) {
+                Py_DECREF(k_key);
+                Py_DECREF(dict);
+                return NULL;
+            }
+            int res = PyDict_SetItem(dict, k_key, v_key);
+            Py_DECREF(k_key);
+            Py_DECREF(v_key);
+            if (res < 0) {
+                Py_DECREF(dict);
+                return NULL;
+            }
+        }
+
+        PyObject *fdict = PyFrozenDict_New(dict);
+        Py_DECREF(dict);
+        if (fdict == NULL) {
+            return NULL;
+        }
+
+        key = _PyTuple_FromPair(fdict, op);
+        Py_DECREF(fdict);
+        return key;
+    }
     else if (PySlice_Check(op)) {
         PySliceObject *slice = (PySliceObject *)op;
         PyObject *start_key = NULL;