From c20e613319ed56a7f349d4d7863d536a24dbd715 Mon Sep 17 00:00:00 2001
From: Eric Traut <eric@traut.com>
Date: Tue, 17 Sep 2024 10:01:45 -0700
Subject: [PATCH] Fixed a bug that results in spurious "variable not allowed in
 type expression" errors when a type alias is used in a loop in both a type
 expression and a value expression. This addresses #9017. (#9018)

---
 .../src/analyzer/codeFlowEngine.ts            | 26 +++++++++++++++++--
 .../src/tests/samples/loop50.py               | 15 +++++++++++
 .../src/tests/typeEvaluator3.test.ts          |  6 +++++
 3 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 packages/pyright-internal/src/tests/samples/loop50.py

diff --git a/packages/pyright-internal/src/analyzer/codeFlowEngine.ts b/packages/pyright-internal/src/analyzer/codeFlowEngine.ts
index 454924357..e32942162 100644
--- a/packages/pyright-internal/src/analyzer/codeFlowEngine.ts
+++ b/packages/pyright-internal/src/analyzer/codeFlowEngine.ts
@@ -401,6 +401,28 @@ export function getCodeFlowEngine(
                 flowNodeTypeCache.cache.delete(flowNode.id);
             }
 
+            // Cleans any "incomplete unknowns" from the specified set of entries
+            // to compute the final type.
+            function cleanIncompleteUnknownForCacheEntry(cacheEntry: FlowNodeTypeResult): Type | undefined {
+                if (!cacheEntry.type) {
+                    return undefined;
+                }
+
+                if (!cacheEntry.incompleteSubtypes || cacheEntry.incompleteSubtypes.length === 0) {
+                    return cleanIncompleteUnknown(cacheEntry.type);
+                }
+
+                const typesToCombine: Type[] = [];
+
+                cacheEntry.incompleteSubtypes?.forEach((entry) => {
+                    if (entry.type && !isIncompleteUnknown(entry.type)) {
+                        typesToCombine.push(cleanIncompleteUnknown(entry.type));
+                    }
+                });
+
+                return combineTypes(typesToCombine);
+            }
+
             function evaluateAssignmentFlowNode(flowNode: FlowAssignment): TypeResult | undefined {
                 // For function and class nodes, the reference node is the name
                 // node, but we need to use the parent node (the FunctionNode or ClassNode)
@@ -456,7 +478,7 @@ export function getCodeFlowEngine(
                         // has changed that may cause the previously-reported incomplete type to change.
                         if (cachedEntry.generationCount === flowIncompleteGeneration) {
                             return FlowNodeTypeResult.create(
-                                cachedEntry.type ? cleanIncompleteUnknown(cachedEntry.type) : undefined,
+                                cleanIncompleteUnknownForCacheEntry(cachedEntry),
                                 /* isIncomplete */ true
                             );
                         }
@@ -966,7 +988,7 @@ export function getCodeFlowEngine(
                     // that have not been evaluated even once, treat it as incomplete. We clean
                     // any incomplete unknowns from the type here to assist with type convergence.
                     return FlowNodeTypeResult.create(
-                        cacheEntry.type ? cleanIncompleteUnknown(cacheEntry.type) : undefined,
+                        cleanIncompleteUnknownForCacheEntry(cacheEntry),
                         /* isIncomplete */ true
                     );
                 }
diff --git a/packages/pyright-internal/src/tests/samples/loop50.py b/packages/pyright-internal/src/tests/samples/loop50.py
new file mode 100644
index 000000000..87b92c36c
--- /dev/null
+++ b/packages/pyright-internal/src/tests/samples/loop50.py
@@ -0,0 +1,15 @@
+# This sample tests the case where a type alias is accessed within
+# a loop as both an annotation and a value expression.
+
+from typing import Literal
+
+
+TA1 = Literal["a", "b"]
+
+
+def func1(values: list):
+    for value in values:
+        x: TA1 = value["x"]
+
+        if x not in TA1.__args__:
+            raise ValueError()
diff --git a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
index 3173de03b..c6ae76913 100644
--- a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
+++ b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
@@ -473,6 +473,12 @@ test('Loop49', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Loop50', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['loop50.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('ForLoop1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['forLoop1.py']);