[Python] Fix async await in branches and cast in closures (#4341)

Author: dbrattli

pyrightconfig.ci.json

@@ -4,7 +4,6 @@
     "**/.venv/**",
     "**/node_modules/**",
     "temp/tests/Python/test_applicative.py",
-    "temp/tests/Python/test_misc.py",
     "temp/tests/Python/fable_modules/thoth_json_python/encode.py"
   ]
 }

src/Fable.Cli/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Fixed
+
+* [Python] Fix `Task<T>` pass-through returns not being awaited in if/else and try/with branches (by @dbrattli)
+* [Python] Fix `:? T as x` type test pattern in closures causing `UnboundLocalError` due to `cast()` shadowing outer variable (by @dbrattli)
+
 ## 5.0.0-alpha.23 - 2026-02-03
 
 ### Changed

src/Fable.Compiler/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Fixed
+
+* [Python] Fix `Task<T>` pass-through returns not being awaited in if/else and try/with branches (by @dbrattli)
+* [Python] Fix `:? T as x` type test pattern in closures causing `UnboundLocalError` due to `cast()` shadowing outer variable (by @dbrattli)
+
 ## 5.0.0-alpha.22 - 2026-02-03
 
 ### Changed

src/Fable.Transforms/Python/Fable2Python.Bases.fs

@@ -149,6 +149,9 @@ let generatePythonProtocolDunders (com: IPythonCompiler) ctx (classEnt: Fable.En
 
     // Generate IMapping dunders: __getitem__, __contains__, __len__, __iter__
     // Note: Method names use Python naming convention (lowercase with underscores)
+    // __iter__ always yields keys following Python's Mapping convention.
+    // F# iteration (for KeyValue(k,v) in map) compiles to GetEnumerator()/MoveNext()/Current
+    // and never uses __iter__, so this is purely for Python interop.
     let mappingDunders =
         if hasIMapping || hasIMutableMapping then
             [
@@ -170,9 +173,7 @@ let generatePythonProtocolDunders (com: IPythonCompiler) ctx (classEnt: Fable.En
                     Arguments.arguments [ Arg.arg "self" ],
                     body = [ Statement.return' (Expression.attribute (self, Identifier "Count", Load)) ]
                 )
-                // def __iter__(self):
-                //     for kv in to_iterator(self.GetEnumerator()):
-                //         yield kv[0]  # kv is a tuple (key, value)
+                // def __iter__(self): yields keys only (Python Mapping convention)
                 let toIterator = com.GetImportExpr(ctx, "fable_library.util", "to_iterator")
                 let kvVar = Expression.name "kv"
 
@@ -184,7 +185,7 @@ let generatePythonProtocolDunders (com: IPythonCompiler) ctx (classEnt: Fable.En
                             Statement.for' (
                                 kvVar,
                                 Expression.call (toIterator, [ selfCall "GetEnumerator" [] ]),
-                                // Access kv[0] since the enumerator yields tuples (key, value)
+                                // Yield kv[0] (key) since GetEnumerator yields (key, value) tuples
                                 [
                                     Statement.expr (
                                         Yield(Some(Expression.subscript (kvVar, Expression.intConstant 0, Load)))

src/Fable.Transforms/Python/Fable2Python.Transforms.fs

@@ -1305,35 +1305,12 @@ let transformTryCatch com (ctx: Context) r returnStrategy (body, catch: option<F
         )
     ]
 
-/// Helper function to generate a cast statement for type narrowing
-let makeCastStatement (com: IPythonCompiler) ctx (ident: Fable.Ident) (typ: Fable.Type) =
-    // Only add cast for generic types where type checker needs help
-    let hasGenerics =
-        match typ with
-        | Fable.DeclaredType(_, genArgs) when not (List.isEmpty genArgs) -> true
-        | Fable.Array _ -> true
-        | Fable.List _ -> true
-        | _ -> false
-
-    // Check if the original type already has the same generic arguments
-    // If so, Pyright can infer the narrowed type and the cast is unnecessary
-    let originalGenArgs = Annotation.getGenericArgs ident.Type
-    let targetGenArgs = Annotation.getGenericArgs typ
-
-    let sameGenericArgs =
-        not (List.isEmpty originalGenArgs)
-        && not (List.isEmpty targetGenArgs)
-        && originalGenArgs = targetGenArgs
-
-    if hasGenerics && not sameGenericArgs then
-        let cast = com.GetImportExpr(ctx, "typing", "cast")
-        let varExpr = identAsExpr com ctx ident
-        let typeAnnotation, importStmts = Annotation.typeAnnotation com ctx None typ
-        let castExpr = Expression.call (cast, [ typeAnnotation; varExpr ])
-        let castStmt = Statement.assign ([ varExpr ], castExpr)
-        importStmts @ [ castStmt ]
-    else
-        []
+/// Helper function to generate a cast statement for type narrowing.
+/// Disabled: The isinstance() check in the if-guard already provides Pyright with
+/// sufficient type narrowing. Generating `value = cast(T, value)` causes Python
+/// UnboundLocalError when inside closures (e.g., task CE bodies), because the
+/// assignment makes Python treat the variable as local throughout the function.
+let makeCastStatement (_com: IPythonCompiler) _ctx (_ident: Fable.Ident) (_typ: Fable.Type) = []
 
 let rec transformIfStatement (com: IPythonCompiler) ctx r ret guardExpr thenStmnt elseStmnt =
     // Create refined context for then/else branches based on guard type

src/Fable.Transforms/Python/Fable2Python.Util.fs

@@ -1125,12 +1125,73 @@ module Helpers =
         let callInfo = Fable.CallInfo.Create(args = [ e ])
         makeIdentExpr "str" |> makeCall None Fable.String callInfo
 
-    /// Transform return statements to wrap their values with await
-    let wrapReturnWithAwait (body: Statement list) : Statement list =
+    /// Transform return statements to wrap their values with await.
+    /// Recursively traverses nested control flow (if/else, match, try, for, while, with)
+    /// to ensure ALL return statements in async functions get awaited.
+    let rec wrapReturnWithAwait (body: Statement list) : Statement list =
         body
         |> List.map (fun stmt ->
             match stmt with
+            | Statement.Return { Value = Some(Await _) } -> stmt // Already awaited
             | Statement.Return { Value = Some value } -> Statement.return' (Await value)
+            | Statement.If ifStmt ->
+                Statement.if' (
+                    ifStmt.Test,
+                    wrapReturnWithAwait ifStmt.Body,
+                    wrapReturnWithAwait ifStmt.Else,
+                    ?loc = ifStmt.Loc
+                )
+            | Statement.Match matchStmt ->
+                let cases =
+                    matchStmt.Cases
+                    |> List.map (fun case ->
+                        MatchCase.matchCase (case.Pattern, wrapReturnWithAwait case.Body, ?guard = case.Guard)
+                    )
+
+                Statement.match' (matchStmt.Subject, cases, ?loc = matchStmt.Loc)
+            | Statement.Try tryStmt ->
+                let handlers =
+                    tryStmt.Handlers
+                    |> List.map (fun h -> { h with Body = wrapReturnWithAwait h.Body })
+
+                Statement.try' (
+                    wrapReturnWithAwait tryStmt.Body,
+                    handlers = handlers,
+                    orElse = wrapReturnWithAwait tryStmt.OrElse,
+                    finalBody = wrapReturnWithAwait tryStmt.FinalBody,
+                    ?loc = tryStmt.Loc
+                )
+            | Statement.For forStmt ->
+                Statement.for' (
+                    forStmt.Target,
+                    forStmt.Iterator,
+                    body = wrapReturnWithAwait forStmt.Body,
+                    orelse = wrapReturnWithAwait forStmt.Else,
+                    ?typeComment = forStmt.TypeComment
+                )
+            | Statement.AsyncFor asyncForStmt ->
+                AsyncFor(
+                    AsyncFor.asyncFor (
+                        asyncForStmt.Target,
+                        asyncForStmt.Iterator,
+                        wrapReturnWithAwait asyncForStmt.Body,
+                        orelse = wrapReturnWithAwait asyncForStmt.Else,
+                        ?typeComment = asyncForStmt.TypeComment
+                    )
+                )
+            | Statement.While whileStmt ->
+                Statement.while' (
+                    whileStmt.Test,
+                    wrapReturnWithAwait whileStmt.Body,
+                    orelse = wrapReturnWithAwait whileStmt.Else,
+                    ?loc = whileStmt.Loc
+                )
+            | Statement.With withStmt ->
+                Statement.with' (
+                    withStmt.Items,
+                    body = wrapReturnWithAwait withStmt.Body,
+                    ?typeComment = withStmt.TypeComment
+                )
             | other -> other
         )
 

tests/Python/TestMap.fs

@@ -333,3 +333,21 @@ let ``test Map works with keys with custom comparison`` () =
     |> Map.add { Bar = "a"; Foo = 10 } 2
     |> Map.count
     |> equal 1
+
+// Note: This compiles to GetEnumerator()/MoveNext()/Current, not __iter__
+[<Fact>]
+let ``test Map iteration with KeyValue yields key-value pairs`` () =
+    let myMap = Map [ "foo", 1; "bar", 2; "baz", 3 ]
+    let mutable keys = []
+    let mutable values = []
+    for KeyValue(key, value) in myMap do
+        keys <- key :: keys
+        values <- value :: values
+    keys |> List.sort |> equal ["bar"; "baz"; "foo"]
+    values |> List.sort |> equal [1; 2; 3]
+
+[<Fact>]
+let ``test Map keys and values work correctly`` () =
+    let myMap = Map [ "a", 10; "b", 20 ] :> IDictionary<_,_>
+    myMap.Keys |> Seq.toList |> equal ["a"; "b"]
+    myMap.Values |> Seq.toList |> equal [10; 20]

tests/Python/TestNonRegression.fs

@@ -284,3 +284,19 @@ let ``test named arguments are converted to snake_case`` () =
     let runner2 = NamedArgsSnakeCase.TestRunner(testCase = "test2", configArgs = [| "arg3" |])
     equal "test2" runner2.TestCase
     equal [| "arg3" |] runner2.ConfigArgs
+
+// Regression: type test pattern (:? T as x) inside closure should not cause
+// UnboundLocalError by reassigning the tested variable
+[<Fact>]
+let ``test type test pattern in closure does not shadow outer variable`` () =
+    let mutable result = ""
+    let processValue (value: obj) =
+        let inner () =
+            match value with
+            | :? string as s -> result <- s
+            | _ -> result <- "not a string"
+        inner ()
+    processValue (box "hello")
+    equal "hello" result
+    processValue (box 42)
+    equal "not a string" result

tests/Python/TestTask.fs

@@ -154,3 +154,48 @@ let ``test TaskCompletionSource is executed correctly`` () =
         x |> fun tsk -> tsk.GetAwaiter().GetResult()
 
     equal 42 result
+
+// Regression: async functions returning Task from if/else branches must await both branches
+let skipPipeline () : Task<int option> = Task.FromResult None
+
+let httpVerb (validate: string -> bool) =
+    fun (value: string) ->
+        if validate value then
+            Task.FromResult(Some 42)
+        else
+            skipPipeline ()
+
+[<Fact>]
+let ``test async pass-through returns are awaited in if branches`` () =
+    let result =
+        httpVerb (fun s -> s = "GET") "GET"
+        |> fun tsk -> tsk.GetAwaiter().GetResult()
+    equal (Some 42) result
+
+    let result2 =
+        httpVerb (fun s -> s = "GET") "POST"
+        |> fun tsk -> tsk.GetAwaiter().GetResult()
+    equal None result2
+
+// Regression: await in nested try/with inside async
+[<Fact>]
+let ``test async returns in try-with are awaited`` () =
+    let getResult (fail: bool) : Task<int> =
+        if fail then
+            failwith "error"
+        else
+            Task.FromResult 99
+
+    let wrapper (fail: bool) =
+        task {
+            try
+                return! getResult fail
+            with
+            | _ -> return -1
+        }
+
+    let result = wrapper false |> fun tsk -> tsk.GetAwaiter().GetResult()
+    equal 99 result
+
+    let result2 = wrapper true |> fun tsk -> tsk.GetAwaiter().GetResult()
+    equal -1 result2