Preserve nested conditional values in comprehension elements

docs/LIMITATIONS.md

@@ -62,8 +62,9 @@ equivalent, or (b) a corner with no occurrence in the 520-module corpus.
 - `with a, b:` renders as nested with statements (equivalent).
 - Async gaps: inlined async comprehensions, async-generator asend/athrow
   protocol, the 3.7 SETUP_EXCEPT-guarded async-for break+else combination.
-- Ternaries inside comprehension elements can be misplaced on some old
-  versions.
+- Nested ternaries in either arm of an inlined comprehension element are
+  reconstructed on 3.12+. Ternaries inside comprehension elements can
+  still be misplaced on older versions, including 3.8–3.11.
 - py2 deep-nesting shape "if containing an if + sibling elif + dangling tail
   return" may lose subsequent branches.
 - A few 3.3 modules (_weakrefset/copy/copyreg) are AST-PASS with an empty

src/decompiler.rs

@@ -51110,10 +51110,16 @@ fn genexpr_ternary_merge(
                 | Op::POP_JUMP_IF_TRUE
                 | Op::POP_JUMP_FORWARD_IF_FALSE
                 | Op::POP_JUMP_FORWARD_IF_TRUE
+                | Op::POP_JUMP_IF_NONE
+                | Op::POP_JUMP_IF_NOT_NONE
+                | Op::POP_JUMP_FORWARD_IF_NONE
+                | Op::POP_JUMP_FORWARD_IF_NOT_NONE
         )
     };
     let mut jf_target = None;
-    for x in instrs.iter().skip(ii + 1) {
+    let mut j = ii + 1;
+    while j < instrs.len() {
+        let x = instrs[j];
         if x.offset >= t {
             break;
         }
@@ -51130,18 +51136,32 @@ fn genexpr_ternary_merge(
             break;
         }
         if is_cj(x.op) && x.target == Some(t) && !x.is_backward {
+            j += 1;
+            continue;
+        }
+        // A then-arm value may itself contain a conditional dictionary or
+        // another ternary. Consume its complete, pure value region before
+        // searching for this arm's final jump, just as the else-arm does.
+        if is_cj(x.op) {
+            let nested_merge = Self::genexpr_ternary_merge(instrs, j)?;
+            if nested_merge > t {
+                return None;
+            }
+            j = instrs.iter().position(|item| item.offset == nested_merge)?;
             continue;
         }
         // 3.12+: TO_BOOL sits between the value run and the cond jump
         // (`x if x % 2 else -x` — without this the scan bailed and the
         // element rendered as a plain `x`, silently dropping the else
         // arm on EVERY version with the folded ternary shape)
         if matches!(x.op, Op::TO_BOOL | Op::NOP | Op::NOT_TAKEN | Op::CACHE) {
+            j += 1;
             continue;
         }
         if !Self::comp_arm_value_op(x.op) {
             return None;
         }
+        j += 1;
     }
     let m = jf_target?;
     let ti = instrs.iter().position(|x| x.offset == t)?;
@@ -51157,13 +51177,7 @@ fn genexpr_ternary_merge(
             j += 1;
             continue;
         }
-        if matches!(
-            x.op,
-            Op::POP_JUMP_IF_FALSE
-                | Op::POP_JUMP_IF_TRUE
-                | Op::POP_JUMP_FORWARD_IF_FALSE
-                | Op::POP_JUMP_FORWARD_IF_TRUE
-        ) {
+        if is_cj(x.op) {
             let nm = Self::genexpr_ternary_merge(instrs, j)?;
             let nmi = instrs.iter().position(|x2| x2.offset == nm)?;
             if nmi > mi {

tests/behavior/cases/b_nested_conditional_elements.py

@@ -0,0 +1,18 @@
+# MIN_VERSION: 3.12
+
+
+def select(rows):
+    return [{
+        "excluded": ({"count": row["count"],
+                      "reasons": loads(row["reasons"]) if row["reasons"] else {}}
+                     if row["count"] is not None else {}),
+        "archived": str(row["archived"]) if row["archived"] else None,
+    } for row in rows]
+
+
+def loads(value):
+    print("loads", value)
+    return {"decoded": value}
+print(select([{"count": 1, "reasons": "why", "archived": 3},
+              {"count": None, "reasons": "skip", "archived": 0},
+              {"count": 2, "reasons": "", "archived": None}]))

tests/fixtures/nested_conditional_elements.py

@@ -0,0 +1,20 @@
+def select(rows):
+    return [{
+        "excluded": ({"count": row["count"],
+                      "reasons": loads(row["reasons"]) if row["reasons"] else {}}
+                     if row["count"] is not None else {}),
+        "archived": str(row["archived"]) if row["archived"] else None,
+    } for row in rows]
+
+
+def then_arm(rows, build):
+    return [build(1 if row.inner else 2) if row.outer else 3 for row in rows]
+
+
+def else_arm(rows, build):
+    return [1 if row.outer else build(2 if row.inner else 3) for row in rows]
+
+
+def both_arms(rows, build):
+    return [build(1 if row.inner is None else 2) if row.outer
+            else build(3 if row.inner is not None else 4) for row in rows]

tests/integration.rs

@@ -169,3 +169,22 @@ fn py314_hello() {
         "py3.14 hello",
     );
 }
+
+#[test]
+fn nested_conditional_comprehension_arms_keep_their_values() {
+    for version in ["3.12", "3.13", "3.14"] {
+        let fixture = format!("nested_conditional_elements.{version}.pyc");
+        let (src, clean) = decompile_fixture(&fixture);
+        assert!(clean, "{fixture}: output marked incomplete:\n{src}");
+        assert_contains(
+            &src,
+            &[
+                "return [{'excluded': {'count': row['count'], 'reasons': loads(row['reasons']) if row['reasons'] else {}} if row['count'] is not None else {}, 'archived': str(row['archived']) if row['archived'] else None} for row in rows]",
+                "return [build(1 if row.inner else 2) if row.outer else 3 for row in rows]",
+                "return [1 if row.outer else build(2 if row.inner else 3) for row in rows]",
+                "return [build(1 if row.inner is None else 2) if row.outer else build(3 if row.inner is not None else 4) for row in rows]",
+            ],
+            &fixture,
+        );
+    }
+}