Keep large dictionary displays inside comprehension elements

src/decompiler.rs

@@ -14797,6 +14797,7 @@ impl<'a> Ctx<'a> {
             | Op::COPY_DICT_WITHOUT_KEYS => {
                 if self.inline_comp.is_some()
                     && matches!(inst.op, Op::LIST_APPEND | Op::SET_ADD | Op::MAP_ADD)
+                    && (!self.version.at_least(3, 10) || arg >= 2)
                 {
                     self.comp_add_element(inst);
                 } else {
@@ -51833,8 +51834,19 @@ fn genexpr_ternary_merge(
                         let v = stack.pop().unwrap_or_else(|| Rc::new(Expr::Name("?".to_string())));
                         (k, v)
                     };
-                    key.get_or_insert(k);
-                    elt.get_or_insert(v);
+                    if self.version.at_least(3, 10) && inst.arg == 1 {
+                        // Large dict displays use MAP_ADD 1 even inside a
+                        // comprehension. They build its element, rather
+                        // than setting the enclosing comprehension's key.
+                        let dictionary = stack.pop()?;
+                        let Expr::Dict(entries) = &*dictionary else { return None; };
+                        let mut entries = entries.clone();
+                        entries.push((k, v));
+                        stack.push(Rc::new(Expr::Dict(entries)));
+                    } else {
+                        key.get_or_insert(k);
+                        elt.get_or_insert(v);
+                    }
                 }
                 Op::YIELD_VALUE => {
                     if let Some(p) = pending_or_filter.take() {

tests/behavior/cases/b_comprehension_dictionary_displays.py

@@ -0,0 +1,25 @@
+# MIN_VERSION: 3.10
+
+
+def list17(rows):
+    return [{
+        "k0": row + 0, "k1": row + 1, "k2": row + 2, "k3": row + 3,
+        "k4": row + 4, "k5": row + 5, "k6": row + 6, "k7": row + 7,
+        "k8": row + 8, "k9": row + 9, "k10": row + 10, "k11": row + 11,
+        "k12": row + 12, "k13": row + 13, "k14": row + 14, "k15": row + 15,
+        "k16": row + 16,
+    } for row in rows]
+
+
+def generator20(rows):
+    return ({
+        "k0": row + 0, "k1": row + 1, "k2": row + 2, "k3": row + 3,
+        "k4": row + 4, "k5": row + 5, "k6": row + 6, "k7": row + 7,
+        "k8": row + 8, "k9": row + 9, "k10": row + 10, "k11": row + 11,
+        "k12": row + 12, "k13": row + 13, "k14": row + 14, "k15": row + 15,
+        "k16": row + 16, "k17": row + 17, "k18": row + 18, "k19": row + 19,
+    } for row in rows)
+
+
+print(list17([0, 10]))
+print(tuple(generator20([1, 20])))

tests/fixtures/comprehension_dictionary_displays.py

@@ -0,0 +1,51 @@
+# Compile with each named CPython version; the committed .pyc files are
+# exercised by the existing Rust integration suite.
+
+# Python 3.10+ uses MAP_ADD 1 for large dictionary displays.
+
+def list17(rows):
+    return [{
+        "k0": row + 0, "k1": row + 1, "k2": row + 2, "k3": row + 3,
+        "k4": row + 4, "k5": row + 5, "k6": row + 6, "k7": row + 7,
+        "k8": row + 8, "k9": row + 9, "k10": row + 10, "k11": row + 11,
+        "k12": row + 12, "k13": row + 13, "k14": row + 14, "k15": row + 15,
+        "k16": row + 16,
+    } for row in rows]
+
+
+def generator17(rows):
+    return ({
+        "k0": row + 0, "k1": row + 1, "k2": row + 2, "k3": row + 3,
+        "k4": row + 4, "k5": row + 5, "k6": row + 6, "k7": row + 7,
+        "k8": row + 8, "k9": row + 9, "k10": row + 10, "k11": row + 11,
+        "k12": row + 12, "k13": row + 13, "k14": row + 14, "k15": row + 15,
+        "k16": row + 16,
+    } for row in rows)
+
+
+def list20(rows):
+    return [{
+        "k0": row + 0, "k1": row + 1, "k2": row + 2, "k3": row + 3,
+        "k4": row + 4, "k5": row + 5, "k6": row + 6, "k7": row + 7,
+        "k8": row + 8, "k9": row + 9, "k10": row + 10, "k11": row + 11,
+        "k12": row + 12, "k13": row + 13, "k14": row + 14, "k15": row + 15,
+        "k16": row + 16, "k17": row + 17, "k18": row + 18, "k19": row + 19,
+    } for row in rows]
+
+
+def generator20(rows):
+    return ({
+        "k0": row + 0, "k1": row + 1, "k2": row + 2, "k3": row + 3,
+        "k4": row + 4, "k5": row + 5, "k6": row + 6, "k7": row + 7,
+        "k8": row + 8, "k9": row + 9, "k10": row + 10, "k11": row + 11,
+        "k12": row + 12, "k13": row + 13, "k14": row + 14, "k15": row + 15,
+        "k16": row + 16, "k17": row + 17, "k18": row + 18, "k19": row + 19,
+    } for row in rows)
+
+
+def dictionary_comprehension(rows):
+    return {row: row + 1 for row in rows}
+
+
+def small_display(rows):
+    return [{"value": row} for row in rows]

tests/integration.rs

@@ -169,3 +169,27 @@ fn py314_hello() {
         "py3.14 hello",
     );
 }
+
+#[test]
+fn large_dictionary_displays_remain_comprehension_elements() {
+    // CPython 3.10 starts using MAP_ADD 1 for large displays. Before 3.12,
+    // list comprehensions also use the separate code-object walker.
+    for version in ["3.10", "3.11", "3.12", "3.13", "3.14"] {
+        let (src, clean) = decompile_fixture(
+            &format!("comprehension_dictionary_displays.{version}.pyc"));
+        assert!(clean, "Python {version}: output marked incomplete:\n{src}");
+        for size in [17, 20] {
+            let fields = (0..size).map(|index| format!("'k{index}': row + {index}"))
+                .collect::<Vec<_>>().join(", ");
+            for (name, open, close) in [("list", "[", "]"), ("generator", "(", ")")] {
+                let expected = format!(
+                    "def {name}{size}(rows):\n    return {open}{{{fields}}} for row in rows{close}");
+                assert_contains(&src, &[&expected], version);
+            }
+        }
+        assert_contains(&src, &[
+            "def dictionary_comprehension(rows):\n    return {row: row + 1 for row in rows}",
+            "def small_display(rows):\n    return [{'value': row} for row in rows]",
+        ], version);
+    }
+}