Fix TypedDict get, pop, and setdefault type evaluation with union keys

packages/pyright-internal/src/analyzer/typeEvaluator.ts

@@ -177,10 +177,12 @@ import { getLastTypedDeclarationForSymbol, isEffectivelyClassVar } from './symbo
 import { assignTupleTypeArgs, expandTuple, getSlicedTupleType, getTypeOfTuple, makeTupleObject } from './tuples';
 import { SpeculativeModeOptions, SpeculativeTypeTracker } from './typeCacheUtils';
 import {
+    applyTypedDictMethodTransform,
     assignToTypedDict,
     assignTypedDictToTypedDict,
     createTypedDictType,
     createTypedDictTypeInlined,
+    getTypedDictClassFromMethod,
     getTypedDictDictEquivalent,
     getTypedDictMappingEquivalent,
     getTypedDictMembersForClass,
@@ -10544,6 +10546,21 @@ export function createTypeEvaluator(
             return { returnType: evaluateCastCall(argList, errorNode) };
         }
 
+        const tdMethodInfo = getTypedDictClassFromMethod(expandedCallType);
+        if (tdMethodInfo) {
+            const tdResult = applyTypedDictMethodTransform(
+                evaluatorInterface,
+                errorNode,
+                argList,
+                tdMethodInfo.classType,
+                tdMethodInfo.methodName,
+                tdMethodInfo.isBound
+            );
+            if (tdResult) {
+                return tdResult;
+            }
+        }
+
         const callResult = validateOverloadedArgTypes(
             errorNode,
             argList,

packages/pyright-internal/src/analyzer/typedDicts.ts

@@ -34,6 +34,7 @@ import { getLastTypedDeclarationForSymbol } from './symbolUtils';
 import {
     Arg,
     AssignTypeFlags,
+    CallResult,
     EvaluatorUsage,
     TypeEvaluator,
     TypeResult,
@@ -53,6 +54,8 @@ import {
     isClassInstance,
     isInstantiableClass,
     isNever,
+    isOverloaded,
+    isUnion,
     maxTypeRecursionCount,
     NeverType,
     OverloadedType,
@@ -1611,6 +1614,214 @@ export function getTypeOfIndexedTypedDict(
     return { type: resultingType, isIncomplete: !!indexTypeResult.isIncomplete };
 }
 
+export function getTypedDictClassFromMethod(
+    type: FunctionType | OverloadedType
+): { classType: ClassType; methodName: string; isBound: boolean } | undefined {
+    const overload = isOverloaded(type) ? OverloadedType.getOverloads(type)[0] : type;
+    if (!overload || !FunctionType.isSynthesizedMethod(overload)) {
+        return undefined;
+    }
+
+    const name = overload.shared?.name;
+    if (name !== 'get' && name !== 'pop' && name !== 'setdefault') {
+        return undefined;
+    }
+
+    let isBound = false;
+    let boundType = overload.priv.strippedFirstParamType;
+    if (boundType) {
+        isBound = true;
+    } else if (overload.shared?.parameters && overload.shared.parameters.length > 0) {
+        boundType = FunctionType.getParamType(overload, 0);
+    }
+
+    if (boundType && isClassInstance(boundType) && ClassType.isTypedDictClass(boundType)) {
+        return { classType: boundType, methodName: name, isBound };
+    }
+
+    return undefined;
+}
+
+export function applyTypedDictMethodTransform(
+    evaluator: TypeEvaluator,
+    errorNode: ExpressionNode,
+    argList: Arg[],
+    typedDictClass: ClassType,
+    methodName: string,
+    isBound: boolean
+): CallResult | undefined {
+    // Validate argument counts and shape:
+    // Bound call: 1 or 2 positional args. Unbound call: 2 or 3 positional args.
+    const minArgs = isBound ? 1 : 2;
+    const maxArgs = isBound ? 2 : 3;
+    if (argList.length < minArgs || argList.length > maxArgs) {
+        return undefined;
+    }
+
+    // Require all arguments to be simple positional (no keyword names, no *args/**kwargs)
+    if (!argList.every((arg) => !arg.name && arg.argCategory === ArgCategory.Simple)) {
+        return undefined;
+    }
+
+    const keyIndex = isBound ? 0 : 1;
+    const defaultIndex = isBound ? 1 : 2;
+
+    // Validate receiver type for unbound method calls
+    if (!isBound) {
+        const selfArg = argList[0];
+        const selfNode = selfArg.valueExpression ?? errorNode;
+        const selfType = (selfArg.typeResult ?? evaluator.getTypeOfExpression(selfNode)).type;
+        const expectedSelfType = ClassType.cloneAsInstance(typedDictClass);
+        if (!evaluator.assignType(expectedSelfType, selfType)) {
+            return undefined;
+        }
+    }
+
+    const keyArg = argList[keyIndex];
+    const keyNode = keyArg.valueExpression ?? errorNode;
+    const keyTypeResult = keyArg.typeResult ?? evaluator.getTypeOfExpression(keyNode);
+    const keyType = keyTypeResult.type;
+
+    if (!isUnion(keyType)) {
+        return undefined;
+    }
+
+    const defaultArg = argList.length > defaultIndex ? argList[defaultIndex] : undefined;
+    const defaultNode = defaultArg?.valueExpression ?? errorNode;
+    const defaultType = defaultArg
+        ? (defaultArg.typeResult ?? evaluator.getTypeOfExpression(defaultNode)).type
+        : undefined;
+
+    const entries = getTypedDictMembersForClass(evaluator, typedDictClass, /* allowNarrowed */ methodName === 'get');
+    let argumentErrors = false;
+
+    const returnType = mapSubtypes(keyType, (keySubtype) => {
+        if (isAnyOrUnknown(keySubtype)) {
+            return keySubtype;
+        }
+
+        if (isClassInstance(keySubtype) && ClassType.isBuiltIn(keySubtype, 'str')) {
+            if (keySubtype.priv.literalValue === undefined) {
+                if (methodName === 'get') {
+                    if (ClassType.isTypedDictEffectivelyClosed(typedDictClass)) {
+                        const extraType = entries.extraItems?.valueType ?? NeverType.createNever();
+                        return defaultType
+                            ? combineTypes([extraType, defaultType])
+                            : combineTypes([extraType, evaluator.getNoneType()]);
+                    }
+                    return defaultType
+                        ? combineTypes([AnyType.create(), defaultType])
+                        : combineTypes([AnyType.create(), evaluator.getNoneType()]);
+                } else if (methodName === 'pop') {
+                    return defaultType ? combineTypes([UnknownType.create(), defaultType]) : UnknownType.create();
+                } else {
+                    return UnknownType.create();
+                }
+            }
+
+            const entryName = keySubtype.priv.literalValue as string;
+            const entry = entries.knownItems.get(entryName) ?? entries.extraItems;
+
+            if (methodName === 'get') {
+                if (entry && !isNever(entry.valueType)) {
+                    if (entry.isRequired || entry.isProvided) {
+                        return entry.valueType;
+                    }
+                    return combineTypes([entry.valueType, defaultType ?? evaluator.getNoneType()]);
+                }
+                if (ClassType.isTypedDictEffectivelyClosed(typedDictClass)) {
+                    const extraType = entries.extraItems?.valueType;
+                    if (extraType) {
+                        return combineTypes([extraType, defaultType ?? evaluator.getNoneType()]);
+                    }
+                    return defaultType ?? evaluator.getNoneType();
+                }
+                return combineTypes([AnyType.create(), defaultType ?? evaluator.getNoneType()]);
+            } else if (methodName === 'pop') {
+                if (entry && !isNever(entry.valueType)) {
+                    if (entry.isReadOnly) {
+                        evaluator.addDiagnostic(
+                            DiagnosticRule.reportTypedDictNotRequiredAccess,
+                            LocAddendum.keyReadOnly().format({
+                                name: entryName,
+                                type: evaluator.printType(typedDictClass),
+                            }),
+                            keyNode
+                        );
+                        argumentErrors = true;
+                        return UnknownType.create();
+                    }
+                    if (entry.isRequired) {
+                        return entry.valueType;
+                    }
+                    return defaultType ? combineTypes([entry.valueType, defaultType]) : entry.valueType;
+                }
+                if (defaultType) {
+                    return defaultType;
+                }
+                evaluator.addDiagnostic(
+                    DiagnosticRule.reportGeneralTypeIssues,
+                    LocAddendum.keyUndefined().format({
+                        name: entryName,
+                        type: evaluator.printType(typedDictClass),
+                    }),
+                    keyNode
+                );
+                argumentErrors = true;
+                return UnknownType.create();
+            } else if (methodName === 'setdefault') {
+                if (entry && !isNever(entry.valueType)) {
+                    if (entry.isReadOnly) {
+                        evaluator.addDiagnostic(
+                            DiagnosticRule.reportGeneralTypeIssues,
+                            LocAddendum.keyReadOnly().format({
+                                name: entryName,
+                                type: evaluator.printType(typedDictClass),
+                            }),
+                            keyNode
+                        );
+                        argumentErrors = true;
+                        return UnknownType.create();
+                    }
+                    if (defaultType && defaultArg) {
+                        const diag = new DiagnosticAddendum();
+                        if (!evaluator.assignType(entry.valueType, defaultType, diag)) {
+                            evaluator.addDiagnostic(
+                                DiagnosticRule.reportArgumentType,
+                                LocMessage.argAssignmentParam().format({
+                                    paramName: 'default',
+                                    paramType: evaluator.printType(entry.valueType),
+                                    argType: evaluator.printType(defaultType),
+                                }) + diag.getString(),
+                                defaultNode
+                            );
+                            argumentErrors = true;
+                        }
+                    }
+                    return entry.valueType;
+                }
+                evaluator.addDiagnostic(
+                    DiagnosticRule.reportGeneralTypeIssues,
+                    LocAddendum.keyUndefined().format({
+                        name: entryName,
+                        type: evaluator.printType(typedDictClass),
+                    }),
+                    keyNode
+                );
+                argumentErrors = true;
+                return UnknownType.create();
+            }
+        }
+
+        return UnknownType.create();
+    });
+
+    return {
+        returnType,
+        argumentErrors,
+    };
+}
+
 // If the specified type has a non-required key, this method marks the
 // key as present.
 export function narrowForKeyAssignment(classType: ClassType, key: string) {

packages/pyright-internal/src/tests/samples/typedDict28.py

@@ -0,0 +1,67 @@
+# This sample tests type inference and diagnostic behavior for TypedDict
+# methods (get, pop, setdefault) when called with union key types.
+
+from typing import Any, Literal, NotRequired, ReadOnly, TypedDict, assert_type, overload
+
+class Person(TypedDict):
+    name: str
+    age: int
+    nickname: NotRequired[str]
+
+class Config(TypedDict):
+    host: ReadOnly[str]
+    port: int
+
+def test_get_union_literal_keys(p: Person, k: Literal["name", "age"]):
+    v = p.get(k)
+    assert_type(v, str | int)
+
+def test_get_union_with_not_required(p: Person, k: Literal["name", "nickname"]):
+    v = p.get(k)
+    assert_type(v, str | None)
+
+def test_get_union_with_unknown_key(p: Person, k: Literal["name", "missing"], default_val: int):
+    v1 = p.get(k)
+    assert_type(v1, str | Any | None)
+
+    v2 = p.get(k, default_val)
+    assert_type(v2, str | Any | int)
+
+def test_pop_union_literal_keys(p: Person, k: Literal["name", "age"]):
+    v = p.pop(k)
+    assert_type(v, str | int)
+
+def test_pop_readonly_diagnostic(c: Config, k: Literal["host", "port"]):
+    # This should report an error because "host" is ReadOnly.
+    c.pop(k)
+
+def test_setdefault_union_literal_keys(p: Person, k: Literal["name", "age"]):
+    # This should report an error for default "val" not matching int ("age").
+    v = p.setdefault(k, "val")
+    assert_type(v, str | int)
+
+def test_unbound_method_union_keys(p: Person, k: Literal["name", "age"]):
+    v = Person.get(p, k)
+    assert_type(v, str | int)
+
+class CustomContainer:
+    @overload
+    def get(self, key: Literal["a"]) -> int: ...
+    @overload
+    def get(self, key: Literal["b"]) -> str: ...
+    def get(self, key: str) -> Any:
+        pass
+
+def test_custom_overloaded_get_not_intercepted(c: CustomContainer, k: Literal["a", "b"]):
+    # User-defined overloaded function should NOT be intercepted by TypedDict transform.
+    # Standard overload resolution should apply.
+    v = c.get(k)
+
+def test_unbound_invalid_receiver(k: Literal["name", "age"]):
+    # Unbound call with invalid receiver should fall back to normal overload validation.
+    Person.get(123, k)
+
+def test_keyword_and_extra_args(p: Person, k: Literal["name", "age"]):
+    # Keyword arguments or extra arguments fall back to standard overload validation.
+    p.get(k, default=0)
+    p.get(k, 0, 1)

packages/pyright-internal/src/tests/typeEvaluator7.test.ts

@@ -808,6 +808,12 @@ test('TypedDict27', () => {
     TestUtils.validateResults(analysisResults, 7);
 });
 
+test('TypedDict28', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typedDict28.py']);
+
+    TestUtils.validateResults(analysisResults, 5);
+});
+
 test('TypedDictInline1', () => {
     const configOptions = new ConfigOptions(Uri.empty());
     configOptions.diagnosticRuleSet.enableExperimentalFeatures = true;