Added check for the illegal use of an await keyword in a lambda. Th…

…is addresses #9406.
microsoft · Nov 7, 2024 · 54b5d0b · 54b5d0b
1 parent 76b81f3
commit 54b5d0b
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/packages/pyright-internal/src/analyzer/binder.ts b/packages/pyright-internal/src/analyzer/binder.ts
@@ -1600,9 +1600,9 @@ export class Binder extends ParseTreeWalker {
 
     override visitAwait(node: AwaitNode) {
         // Make sure this is within an async lambda or function.
-        const enclosingFunction = ParseTreeUtils.getEnclosingFunction(node);
-        if (enclosingFunction === undefined || !enclosingFunction.d.isAsync) {
-            if (this._fileInfo.ipythonMode && enclosingFunction === undefined) {
+        const execScopeNode = ParseTreeUtils.getExecutionScopeNode(node);
+        if (execScopeNode?.nodeType !== ParseNodeType.Function || !execScopeNode.d.isAsync) {
+            if (this._fileInfo.ipythonMode && execScopeNode?.nodeType === ParseNodeType.Module) {
                 // Top level await is allowed in ipython mode.
                 return true;
             }

diff --git a/packages/pyright-internal/src/tests/samples/coroutines1.py b/packages/pyright-internal/src/tests/samples/coroutines1.py
@@ -16,6 +16,15 @@ async def coroutine1():
 await a
 
 
+async def func1() -> int: ...
+
+
+async def func2() -> None:
+    # This should generate an error because await cannot be
+    # used in a lambda.
+    x = lambda: await func2()
+
+
 def needs_int(val: int):
     pass
 

diff --git a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
@@ -159,7 +159,7 @@ test('Coroutines1', () => {
     configOptions.defaultPythonVersion = pythonVersion3_10;
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['coroutines1.py'], configOptions);
 
-    TestUtils.validateResults(analysisResults, 4);
+    TestUtils.validateResults(analysisResults, 5);
 });
 
 test('Coroutines2', () => {