Added missing checks for illegal use of await and async within li…

…st, set and dictionary comprehensions within a non-async function. This addresses #9414.
microsoft · Nov 7, 2024 · 3b32af0 · 3b32af0
1 parent ff3902f
commit 3b32af0
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 6 deletions.
diff --git a/packages/pyright-internal/src/analyzer/binder.ts b/packages/pyright-internal/src/analyzer/binder.ts
@@ -1607,11 +1607,17 @@ export class Binder extends ParseTreeWalker {
                 return true;
             }
 
+            const isInGenerator =
+                node.parent?.nodeType === ParseNodeType.Comprehension &&
+                node.parent?.parent?.nodeType !== ParseNodeType.List &&
+                node.parent?.parent?.nodeType !== ParseNodeType.Set &&
+                node.parent?.parent?.nodeType !== ParseNodeType.Dictionary;
+
             // Allow if it's within a generator expression. Execution of
             // generator expressions is deferred and therefore can be
             // run within the context of an async function later.
-            if (node.parent?.nodeType !== ParseNodeType.Comprehension) {
-                this._addSyntaxError(LocMessage.awaitNotInAsync(), node);
+            if (!isInGenerator) {
+                this._addSyntaxError(LocMessage.awaitNotInAsync(), node.d.awaitToken);
             }
         }
 
@@ -2181,7 +2187,11 @@ export class Binder extends ParseTreeWalker {
                                 // Allow if it's within a generator expression. Execution of
                                 // generator expressions is deferred and therefore can be
                                 // run within the context of an async function later.
-                                if (node.parent?.nodeType === ParseNodeType.List) {
+                                if (
+                                    node.parent?.nodeType === ParseNodeType.List ||
+                                    node.parent?.nodeType === ParseNodeType.Set ||
+                                    node.parent?.nodeType === ParseNodeType.Dictionary
+                                ) {
                                     this._addSyntaxError(LocMessage.asyncNotInAsyncFunction(), compr.d.asyncToken);
                                 }
                             }

diff --git a/packages/pyright-internal/src/parser/parseNodes.ts b/packages/pyright-internal/src/parser/parseNodes.ts
@@ -1262,7 +1262,8 @@ export namespace AugmentedAssignmentNode {
 export interface AwaitNode extends ParseNodeBase<ParseNodeType.Await> {
     d: {
         expr: ExpressionNode;
-        hasParens?: boolean;
+        awaitToken: Token;
+        hasParens: boolean;
     };
 }
 
@@ -1275,7 +1276,7 @@ export namespace AwaitNode {
             id: _nextNodeId++,
             parent: undefined,
             a: undefined,
-            d: { expr },
+            d: { expr, awaitToken, hasParens: false },
         };
 
         expr.parent = node;
@@ -1418,7 +1419,7 @@ export interface ComprehensionNode extends ParseNodeBase<ParseNodeType.Comprehen
         expr: ParseNode;
         forIfNodes: ComprehensionForIfNode[];
         isGenerator: boolean;
-        hasParens?: boolean;
+        hasParens: boolean;
     };
 }
 
@@ -1435,6 +1436,7 @@ export namespace ComprehensionNode {
                 expr,
                 forIfNodes: [],
                 isGenerator,
+                hasParens: false,
             },
         };
 

diff --git a/packages/pyright-internal/src/tests/samples/await3.py b/packages/pyright-internal/src/tests/samples/await3.py
@@ -0,0 +1,37 @@
+# This sample tests various places where await is invalid.
+
+from typing import Any
+
+
+def func1() -> Any: ...
+
+
+def func2():
+    # These are OK because generators can be called
+    # outside of the context of the current function.
+    (v async for v in func1())
+    (await v for v in func1())
+
+    # This should generate an error because async
+    # cannot be used outside of an async function.
+    [x async for x in func1()]
+
+    # This should generate an error because async
+    # cannot be used outside of an async function.
+    {x async for x in func1()}
+
+    # This should generate an error because async
+    # cannot be used outside of an async function.
+    {k: v async for k, v in func1()}
+
+    # This should generate an error because await
+    # cannot be used outside of an async function.
+    (x for x in await func1())
+
+    # This should generate an error because await
+    # cannot be used outside of an async function.
+    [await x for x in func1()]
+
+    # This should generate an error because await
+    # cannot be used outside of an async function.
+    {await k: v for k, v in func1()}
diff --git a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
@@ -151,6 +151,12 @@ test('Await2', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Await3', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['await3.py']);
+
+    TestUtils.validateResults(analysisResults, 6);
+});
+
 test('Coroutines1', () => {
     const configOptions = new ConfigOptions(Uri.empty());