From 4d8b903c80bb456824746e20dcaa96336afd0dc5 Mon Sep 17 00:00:00 2001
From: Eric Traut <eric@traut.com>
Date: Mon, 13 Jan 2025 18:58:36 -0700
Subject: [PATCH] Fixed bug that results in a spurious error when an assignment
 expression (walrus operator) is used as an argument to a constructor call
 when the constructor has both a `__new__` and `__init__` with differing
 bidirectional type inference contexts. This addresses #9659. (#9699)

---
 .../src/analyzer/constructors.ts              | 24 +++++++++++--------
 .../src/analyzer/typeCacheUtils.ts            |  2 +-
 .../src/analyzer/typeEvaluatorTypes.ts        |  7 +++++-
 .../src/tests/samples/constructor33.py        | 14 +++++++++++
 .../src/tests/typeEvaluator6.test.ts          |  6 +++++
 5 files changed, 41 insertions(+), 12 deletions(-)
 create mode 100644 packages/pyright-internal/src/tests/samples/constructor33.py

diff --git a/packages/pyright-internal/src/analyzer/constructors.ts b/packages/pyright-internal/src/analyzer/constructors.ts
index effdb91f1e4b..c334c4ec6992 100644
--- a/packages/pyright-internal/src/analyzer/constructors.ts
+++ b/packages/pyright-internal/src/analyzer/constructors.ts
@@ -423,16 +423,20 @@ function validateNewMethod(
 
     const constraints = new ConstraintTracker();
 
-    const callResult = evaluator.useSpeculativeMode(useSpeculativeModeForArgs ? errorNode : undefined, () => {
-        return evaluator.validateCallArgs(
-            errorNode,
-            argList,
-            newMethodTypeResult,
-            constraints,
-            skipUnknownArgCheck,
-            inferenceContext
-        );
-    });
+    const callResult = evaluator.useSpeculativeMode(
+        useSpeculativeModeForArgs ? errorNode : undefined,
+        () => {
+            return evaluator.validateCallArgs(
+                errorNode,
+                argList,
+                newMethodTypeResult,
+                constraints,
+                skipUnknownArgCheck,
+                inferenceContext
+            );
+        },
+        { dependentType: newMethodTypeResult.type }
+    );
 
     if (callResult.isTypeIncomplete) {
         isTypeIncomplete = true;
diff --git a/packages/pyright-internal/src/analyzer/typeCacheUtils.ts b/packages/pyright-internal/src/analyzer/typeCacheUtils.ts
index e6b26a0c1e9b..c893f8cff83b 100644
--- a/packages/pyright-internal/src/analyzer/typeCacheUtils.ts
+++ b/packages/pyright-internal/src/analyzer/typeCacheUtils.ts
@@ -58,7 +58,7 @@ export interface SpeculativeModeOptions {
 // a context is popped off the stack, all of the speculative type cache
 // entries that were created within that context are removed from the
 // corresponding type caches because they are no longer valid.
-// Each type context also contains a map of "speculative types" that are
+// The tracker also also contains a map of "speculative types" that are
 // contextually evaluated based on an "expected type" and potentially
 // one or more "dependent types". The "expected type" applies in cases
 // where the speculative root node is being evaluated with bidirectional
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts b/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts
index 69fc859642b4..ba6f4b2d6c04 100644
--- a/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts
+++ b/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts
@@ -36,6 +36,7 @@ import { Declaration } from './declaration';
 import { ResolvedAliasInfo } from './declarationUtils';
 import { SymbolWithScope } from './scope';
 import { Symbol, SynthesizedTypeInfo } from './symbol';
+import { SpeculativeModeOptions } from './typeCacheUtils';
 import { PrintTypeFlags } from './typePrinter';
 import {
     AnyType,
@@ -843,7 +844,11 @@ export interface TypeEvaluator {
 
     getTypeCacheEntryCount: () => number;
     disposeEvaluator: () => void;
-    useSpeculativeMode: <T>(speculativeNode: ParseNode | undefined, callback: () => T) => T;
+    useSpeculativeMode: <T>(
+        speculativeNode: ParseNode | undefined,
+        callback: () => T,
+        options?: SpeculativeModeOptions
+    ) => T;
     isSpeculativeModeInUse: (node: ParseNode | undefined) => boolean;
     setTypeResultForNode: (node: ParseNode, typeResult: TypeResult, flags?: EvalFlags) => void;
 
diff --git a/packages/pyright-internal/src/tests/samples/constructor33.py b/packages/pyright-internal/src/tests/samples/constructor33.py
new file mode 100644
index 000000000000..3d0bce93e997
--- /dev/null
+++ b/packages/pyright-internal/src/tests/samples/constructor33.py
@@ -0,0 +1,14 @@
+# This sample tests the case where an argument to a constructor
+# uses an assignment expression (walrus operator) and the constructor
+# has both a __new__ and __init__ method whose parameters have
+# different bidirectional type inference contexts.
+
+from typing import Any, Self
+
+
+class A:
+    def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...
+    def __init__(self, base: list[str], joined: str) -> None: ...
+
+
+A(temp := ["x"], " ".join(temp))
diff --git a/packages/pyright-internal/src/tests/typeEvaluator6.test.ts b/packages/pyright-internal/src/tests/typeEvaluator6.test.ts
index 491b21a5acb0..0e5e7f28e947 100644
--- a/packages/pyright-internal/src/tests/typeEvaluator6.test.ts
+++ b/packages/pyright-internal/src/tests/typeEvaluator6.test.ts
@@ -849,6 +849,12 @@ test('Constructor32', () => {
     TestUtils.validateResults(analysisResults, 1);
 });
 
+test('Constructor33', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['constructor33.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('ConstructorCallable1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['constructorCallable1.py']);