generic ? is the wrapped optional

Author: KotlinIsland

crates/by_transforms/src/transforms/coalesce.rs

@@ -12,20 +12,22 @@ use crate::type_info::TypeInfo;
 /// inside the rewrite instead of being clobbered by first-wins overlap dedup
 pub(crate) struct NoneCoalesce<'src> {
     source: &'src str,
+    types: &'src dyn TypeInfo,
     pub(crate) edits: Vec<(TextRange, Vec<Fragment>)>,
 }
 
 impl<'src> NoneCoalesce<'src> {
-    pub(crate) fn new(source: &'src str) -> Self {
+    pub(crate) fn new(source: &'src str, types: &'src dyn TypeInfo) -> Self {
         Self {
             source,
+            types,
             edits: Vec::new(),
         }
     }
 }
 
-fn expand_none_chain(expr: &Expr, source: &str) -> Option<String> {
-    let (form, guards) = super::none_chain::expand_chain(expr, source)?;
+fn expand_none_chain(expr: &Expr, source: &str, types: &dyn TypeInfo) -> Option<String> {
+    let (form, guards) = super::none_chain::expand_chain(expr, source, types)?;
     Some(super::none_chain::build_expansion(&guards, &form, "_t"))
 }
 
@@ -68,10 +70,17 @@ impl NoneCoalesce<'_> {
             return self.operand_value(&b.left);
         }
         let rhs = self.operand_value(&b.right);
-        match expand_none_chain(&b.left, self.source) {
+        // a wrapped optional (`int??`, a generic `T?`) keeps its present value
+        // inside the runtime wrapper — the present branch unwraps with `.value`
+        let unwrap = if self.types.wrapped_optional(&b.left) {
+            ".value"
+        } else {
+            ""
+        };
+        match expand_none_chain(&b.left, self.source, self.types) {
             Some(expanded) => {
                 let mut t = vec![
-                    Fragment::Lit("_t if (_t := ".to_owned()),
+                    Fragment::Lit(format!("_t{unwrap} if (_t := ")),
                     Fragment::Lit(expanded),
                     Fragment::Lit(") is not None else ".to_owned()),
                 ];
@@ -81,7 +90,7 @@ impl NoneCoalesce<'_> {
             None if is_trivially_pure(&b.left) => {
                 let mut t = vec![
                     Fragment::Src(b.left.range()),
-                    Fragment::Lit(" if ".to_owned()),
+                    Fragment::Lit(format!("{unwrap} if ")),
                     Fragment::Src(b.left.range()),
                     Fragment::Lit(" is not None else ".to_owned()),
                 ];
@@ -92,7 +101,7 @@ impl NoneCoalesce<'_> {
                 // a chained (left-associative) `??` puts another coalesce on the
                 // left; recurse through `operand_value` so it is lowered rather
                 // than copied verbatim. hoist to walrus so the LHS runs once
-                let mut t = vec![Fragment::Lit("_t if (_t := ".to_owned())];
+                let mut t = vec![Fragment::Lit(format!("_t{unwrap} if (_t := "))];
                 t.extend(self.operand_value(&b.left));
                 t.push(Fragment::Lit(") is not None else ".to_owned()));
                 t.extend(rhs);
@@ -110,7 +119,7 @@ impl NoneCoalesce<'_> {
         {
             return self.expand_coalesce(expr);
         }
-        if let Some(expanded) = expand_none_chain(expr, self.source) {
+        if let Some(expanded) = expand_none_chain(expr, self.source, self.types) {
             return vec![Fragment::Lit(expanded)];
         }
         vec![Fragment::Src(expr.range())]
@@ -161,8 +170,8 @@ impl<'src> NoneCoalescePass<'src> {
 }
 
 impl TypeAwarePass for NoneCoalescePass<'_> {
-    fn run(&self, stmts: &[Stmt], _types: &dyn TypeInfo, ctx: &mut PassContext) {
-        let mut inner = NoneCoalesce::new(self.source);
+    fn run(&self, stmts: &[Stmt], types: &dyn TypeInfo, ctx: &mut PassContext) {
+        let mut inner = NoneCoalesce::new(self.source, types);
         for stmt in stmts {
             inner.visit_stmt(stmt);
         }
@@ -254,6 +263,21 @@ mod tests {
         );
     }
 
+    #[test]
+    fn wrapped_lhs_unwraps_present_value() {
+        // a wrapped optional's present value lives inside the runtime wrapper —
+        // the present branch reads `.value`
+        let out = transpile(
+            "def g() -> int??:\n    return Some(5)\nx = g() ?? -1\n",
+            &Config::test_default(),
+        )
+        .unwrap();
+        assert!(
+            out.contains("x = _t.value if (_t := g()) is not None else -1\n"),
+            "got: {out}"
+        );
+    }
+
     #[test]
     fn python_unchanged() {
         let src = "x = a if a is not None else b\n";

crates/by_transforms/src/transforms/none_chain.rs

@@ -40,18 +40,31 @@ fn pick_temp_var(types: &dyn TypeInfo, anchor: &Expr) -> &'static str {
 /// walks an attribute-access chain and returns `Some((python_form, guards))` when
 /// any `?.` is present, where `python_form` has all `?.` replaced by `.` and
 /// `guards` is the ordered list of accumulated sub-expressions that must be
-/// non-None before each subsequent optional access is safe
-pub(super) fn expand_chain(expr: &Expr, source: &str) -> Option<(String, Vec<String>)> {
+/// non-None before each subsequent optional access is safe.
+///
+/// a `?.` whose base is a *wrapped* optional (`int??`, a generic `T?`)
+/// reaches its present value through the runtime wrapper: the guard still
+/// tests the wrapper against `None`, but the access reads `.value` first
+pub(super) fn expand_chain(
+    expr: &Expr,
+    source: &str,
+    types: &dyn TypeInfo,
+) -> Option<(String, Vec<String>)> {
     let Expr::Attribute(attr) = expr else {
         return None;
     };
     let field = attr.attr.as_str();
-    match expand_chain(&attr.value, source) {
+    let unwrap = if attr.optional && types.wrapped_optional(&attr.value) {
+        ".value"
+    } else {
+        ""
+    };
+    match expand_chain(&attr.value, source, types) {
         Some((v_form, mut guards)) => {
             if attr.optional {
                 guards.push(v_form.clone());
             }
-            Some((format!("{v_form}.{field}"), guards))
+            Some((format!("{v_form}{unwrap}.{field}"), guards))
         }
         None => {
             if !attr.optional {
@@ -60,7 +73,7 @@ pub(super) fn expand_chain(expr: &Expr, source: &str) -> Option<(String, Vec<Str
             let start = usize::from(attr.value.range().start());
             let end = usize::from(attr.value.range().end());
             let v_form = source[start..end].to_owned();
-            Some((format!("{v_form}.{field}"), vec![v_form]))
+            Some((format!("{v_form}{unwrap}.{field}"), vec![v_form]))
         }
     }
 }
@@ -132,7 +145,7 @@ impl<'ast> Visitor<'ast> for NoneChain<'_> {
 
     fn visit_expr(&mut self, expr: &'ast Expr) {
         if let Expr::Attribute(_) = expr {
-            if let Some((form, guards)) = expand_chain(expr, self.source) {
+            if let Some((form, guards)) = expand_chain(expr, self.source, self.types) {
                 let temp = pick_temp_var(self.types, expr);
                 self.edits.push(Fix::safe_edit(Edit::range_replacement(
                     build_expansion(&guards, &form, temp),
@@ -157,6 +170,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn wrapped_base_unwraps_value() {
+        // `?.` on a wrapped optional reads the present value through the
+        // runtime wrapper: the guard tests the wrapper, the access goes
+        // through `.value`
+        let out = transpile(
+            "def g() -> int??:\n    return Some(5)\nw = g()\nx = w?.bit_length\n",
+            &Config::test_default(),
+        )
+        .unwrap();
+        assert!(
+            out.contains("x = None if w is None else w.value.bit_length\n"),
+            "got: {out}"
+        );
+    }
+
     #[test]
     fn basic_chain() {
         check("x = a?.b\n", "x = None if a is None else a.b\n");

crates/by_transforms/src/transforms/optional_type.rs

@@ -36,6 +36,11 @@ struct OptionalLower<'src> {
     /// set when any lowered optional produced a runtime `Optional[...]` wrapper
     needs_runtime: bool,
     source: &'src str,
+    /// stack of in-scope PEP 695 type-parameter names. `?` over a bare type
+    /// variable lowers to the *wrapped* form (`Optional[T | None]`) — a plain
+    /// union would flatten when `T` binds to an optional (mirrors ty's typing
+    /// of a generic `T?`)
+    typevar_scopes: Vec<Vec<String>>,
 }
 
 impl<'src> OptionalLower<'src> {
@@ -44,13 +49,39 @@ impl<'src> OptionalLower<'src> {
             edits: Vec::new(),
             needs_runtime: false,
             source,
+            typevar_scopes: Vec::new(),
         }
     }
+
+    fn in_scope_typevar(&self, name: &str) -> bool {
+        self.typevar_scopes
+            .iter()
+            .any(|scope| scope.iter().any(|n| n == name))
+    }
 }
 
 impl<'ast> Visitor<'ast> for OptionalLower<'_> {
     fn visit_stmt(&mut self, stmt: &'ast Stmt) {
+        let type_params = match stmt {
+            Stmt::FunctionDef(f) => f.type_params.as_deref(),
+            Stmt::ClassDef(c) => c.type_params.as_deref(),
+            _ => None,
+        };
+        let pushed = if let Some(tp) = type_params {
+            self.typevar_scopes.push(
+                tp.type_params
+                    .iter()
+                    .map(|p| p.name().id.as_str().to_owned())
+                    .collect(),
+            );
+            true
+        } else {
+            false
+        };
         walk_stmt(self, stmt);
+        if pushed {
+            self.typevar_scopes.pop();
+        }
     }
 
     fn visit_expr(&mut self, expr: &'ast Expr) {
@@ -82,16 +113,21 @@ impl<'ast> Visitor<'ast> for OptionalLower<'_> {
         let node = unary.range;
         let tail = &self.source[usize::from(inner.range().end())..usize::from(node.end())];
         let close_parens: String = tail.chars().filter(|c| *c == ')').collect();
-        if depth >= 2 {
+        // a bare in-scope type variable gets one extra wrapper layer: `T?` is
+        // `Optional[T | None]`, matching ty's wrapped typing of a generic `?`
+        let generic_operand =
+            matches!(inner, Expr::Name(n) if self.in_scope_typevar(n.id.as_str()));
+        let wrap_layers = (depth - 1) as usize + usize::from(generic_operand);
+        if wrap_layers >= 1 {
             self.needs_runtime = true;
             self.edits.push((
                 TextRange::empty(node.start()),
-                "Optional[".repeat((depth - 1) as usize),
+                "Optional[".repeat(wrap_layers),
             ));
         }
         let mut replacement = close_parens;
         replacement.push_str(" | None");
-        for _ in 1..depth {
+        for _ in 0..wrap_layers {
             replacement.push(']');
         }
         self.edits
@@ -210,6 +246,32 @@ class Optional:
         );
     }
 
+    #[test]
+    fn generic_typevar_optional_wraps() {
+        // `?` over a bare in-scope type variable is the wrapped form — a plain
+        // union would flatten when `T` binds to an optional. (the 3.10 generics
+        // polyfill's `T` → `_T` rename composes inside the wrapper)
+        check(
+            indoc! {"
+                def f[T](t: T) -> T?:
+                    return Some(t)
+            "},
+            &format!(
+                "from typing import TypeVar\n{RUNTIME}_T = TypeVar(\"_T\")\ndef f(t: _T) -> Optional[_T | None]:\n    return Optional(t)\n"
+            ),
+        );
+    }
+
+    #[test]
+    fn non_typevar_name_optional_stays_union() {
+        // a name that is not an in-scope type parameter lowers to the plain
+        // union as before
+        check(
+            "def f(x: int?) -> None: ...\n",
+            "def f(x: int | None) -> None: ...\n",
+        );
+    }
+
     #[test]
     fn optional_of_negation_composes() {
         // the operand source is preserved, so the `not_type` transform still

crates/by_transforms/src/transforms/propagate.rs

@@ -34,8 +34,8 @@ use crate::type_info::{AbsentTest, TypeInfo};
 /// result-like `T | E` tests `isinstance(_, BaseException)`.
 fn absent_condition(test: AbsentTest, target: &str) -> String {
     match test {
-        AbsentTest::IsNone => format!("{target} is None"),
-        AbsentTest::IsException => format!("isinstance({target}, BaseException)"),
+        AbsentTest::Optional | AbsentTest::WrappedOptional => format!("{target} is None"),
+        AbsentTest::Result => format!("isinstance({target}, BaseException)"),
     }
 }
 
@@ -168,21 +168,27 @@ impl<'ast> Visitor<'ast> for Propagate<'_> {
             let absent = self
                 .types
                 .propagate_absent_test(&unary.operand)
-                .unwrap_or(AbsentTest::IsNone);
+                .unwrap_or(AbsentTest::Optional);
             let operand_src = self.src(unary.operand.range()).to_owned();
+            // a wrapped optional's present value is inside the runtime wrapper
+            let unwrap = if absent == AbsentTest::WrappedOptional {
+                ".value"
+            } else {
+                ""
+            };
             let (guard, value) = if is_trivially_pure(&unary.operand) {
                 let cond = absent_condition(absent, &operand_src);
                 (
                     format!("if {cond}: return {operand_src}\n{indent}"),
-                    operand_src,
+                    format!("{operand_src}{unwrap}"),
                 )
             } else {
                 let temp = format!("_prop{}", self.counter);
                 self.counter += 1;
                 let cond = absent_condition(absent, &temp);
                 (
                     format!("{temp} = {operand_src}\n{indent}if {cond}: return {temp}\n{indent}"),
-                    temp,
+                    format!("{temp}{unwrap}"),
                 )
             };
             // hoist the guard immediately before the enclosing statement, then
@@ -235,6 +241,42 @@ mod tests {
         assert!(err.contains(needle), "got: {err}");
     }
 
+    #[test]
+    fn wrapped_operand_unwraps_value() {
+        // a wrapped optional's present value lives inside the runtime wrapper:
+        // the guard returns the absent `None`, the use site reads `.value`
+        check(
+            indoc::indoc! {"
+                def g() -> int??:
+                    return Some(5)
+
+                def f() -> int?:
+                    x = g()^
+                    return x
+            "},
+            indoc::indoc! {"
+                class Optional:
+                    def __init__(self, value):
+                        self.value = value
+
+                    def __class_getitem__(cls, item):
+                        return cls
+
+                    def __repr__(self):
+                        return f\"Some({self.value!r})\"
+
+                def g() -> Optional[int | None]:
+                    return Optional(5)
+
+                def f() -> int | None:
+                    _prop0 = g()
+                    if _prop0 is None: return _prop0
+                    x = _prop0.value
+                    return x
+            "},
+        );
+    }
+
     #[test]
     fn propagate_outside_function_is_an_error() {
         // hoisting a `return` to module scope, into a lambda, or out of a