From ee786485d22161b0bc97efcb75c6d9169ec5a140 Mon Sep 17 00:00:00 2001
From: Eric Traut <eric@traut.com>
Date: Fri, 29 Nov 2024 19:13:05 +0100
Subject: [PATCH] Fixed a bug that leads to a false positive when a function is
 decorated and has no explicit return type annotation and the body references
 the decorated function in a loop. This addresses #9492.

---
 .../pyright-internal/src/analyzer/decorators.ts     | 12 +++++++++++-
 .../pyright-internal/src/analyzer/typeEvaluator.ts  |  5 ++++-
 .../pyright-internal/src/tests/samples/loop52.py    | 13 +++++++++++++
 .../src/tests/typeEvaluator3.test.ts                |  6 ++++++
 4 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 packages/pyright-internal/src/tests/samples/loop52.py

diff --git a/packages/pyright-internal/src/analyzer/decorators.ts b/packages/pyright-internal/src/analyzer/decorators.ts
index ce69cf920c16..862fe6b0f043 100644
--- a/packages/pyright-internal/src/analyzer/decorators.ts
+++ b/packages/pyright-internal/src/analyzer/decorators.ts
@@ -176,7 +176,17 @@ export function applyFunctionDecorator(
         }
     }
 
-    let returnType = getTypeOfDecorator(evaluator, decoratorNode, inputFunctionType);
+    // Clear the PartiallyEvaluated flag in the input if it's set so
+    // it doesn't propagate to the decorated type.
+    const decoratorArg =
+        isFunction(inputFunctionType) && FunctionType.isPartiallyEvaluated(inputFunctionType)
+            ? FunctionType.cloneWithNewFlags(
+                  inputFunctionType,
+                  inputFunctionType.shared.flags & ~FunctionTypeFlags.PartiallyEvaluated
+              )
+            : inputFunctionType;
+
+    let returnType = getTypeOfDecorator(evaluator, decoratorNode, decoratorArg);
 
     // Check for some built-in decorator types with known semantics.
     if (isFunction(decoratorType)) {
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
index a5987c04fae9..7a17f7915f82 100644
--- a/packages/pyright-internal/src/analyzer/typeEvaluator.ts
+++ b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -18354,7 +18354,6 @@ export function createTypeEvaluator(
         // Set the "partially evaluated" flag around this logic to detect recursion.
         functionType.shared.flags |= FunctionTypeFlags.PartiallyEvaluated;
         const preDecoratedType = node.d.isAsync ? createAsyncFunction(node, functionType) : functionType;
-        functionType.shared.flags &= ~FunctionTypeFlags.PartiallyEvaluated;
 
         // Apply all of the decorators in reverse order.
         decoratedType = preDecoratedType;
@@ -18402,6 +18401,10 @@ export function createTypeEvaluator(
 
         writeTypeCache(node, { type: decoratedType }, EvalFlags.None);
 
+        // Now that the decorator has been applied, we can clear the
+        // "partially evaluated" flag.
+        functionType.shared.flags &= ~FunctionTypeFlags.PartiallyEvaluated;
+
         return { functionType, decoratedType };
     }
 
diff --git a/packages/pyright-internal/src/tests/samples/loop52.py b/packages/pyright-internal/src/tests/samples/loop52.py
new file mode 100644
index 000000000000..c1a4c87264d2
--- /dev/null
+++ b/packages/pyright-internal/src/tests/samples/loop52.py
@@ -0,0 +1,13 @@
+# This sample tests the case where a function accesses its own decorated
+# form within a loop.
+
+from contextlib import contextmanager
+
+
+@contextmanager
+def func1():
+    yield
+
+    for _ in ():
+        with func1():
+            return
diff --git a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
index 55f9b9763901..6d08c63de73e 100644
--- a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
+++ b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts
@@ -497,6 +497,12 @@ test('Loop51', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Loop52', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['loop52.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('ForLoop1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['forLoop1.py']);