diff --git a/src/decompiler.rs b/src/decompiler.rs index 33a3465d..bce2fb97 100644 --- a/src/decompiler.rs +++ b/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() { diff --git a/tests/behavior/cases/b_comprehension_dictionary_displays.py b/tests/behavior/cases/b_comprehension_dictionary_displays.py new file mode 100644 index 00000000..68c35a62 --- /dev/null +++ b/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]))) diff --git a/tests/fixtures/comprehension_dictionary_displays.py b/tests/fixtures/comprehension_dictionary_displays.py new file mode 100644 index 00000000..edf421b6 --- /dev/null +++ b/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] diff --git a/tests/integration.rs b/tests/integration.rs index f4af5082..b4ba9c03 100644 --- a/tests/integration.rs +++ b/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::>().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); + } +} diff --git a/tests/pyc/comprehension_dictionary_displays.3.10.pyc b/tests/pyc/comprehension_dictionary_displays.3.10.pyc new file mode 100644 index 00000000..f7587b2e Binary files /dev/null and b/tests/pyc/comprehension_dictionary_displays.3.10.pyc differ diff --git a/tests/pyc/comprehension_dictionary_displays.3.11.pyc b/tests/pyc/comprehension_dictionary_displays.3.11.pyc new file mode 100644 index 00000000..e0866250 Binary files /dev/null and b/tests/pyc/comprehension_dictionary_displays.3.11.pyc differ diff --git a/tests/pyc/comprehension_dictionary_displays.3.12.pyc b/tests/pyc/comprehension_dictionary_displays.3.12.pyc new file mode 100644 index 00000000..dbdc5f36 Binary files /dev/null and b/tests/pyc/comprehension_dictionary_displays.3.12.pyc differ diff --git a/tests/pyc/comprehension_dictionary_displays.3.13.pyc b/tests/pyc/comprehension_dictionary_displays.3.13.pyc new file mode 100644 index 00000000..c1494d26 Binary files /dev/null and b/tests/pyc/comprehension_dictionary_displays.3.13.pyc differ diff --git a/tests/pyc/comprehension_dictionary_displays.3.14.pyc b/tests/pyc/comprehension_dictionary_displays.3.14.pyc new file mode 100644 index 00000000..340d8546 Binary files /dev/null and b/tests/pyc/comprehension_dictionary_displays.3.14.pyc differ