Added support for dictionary expansion of a TypedDict within a dictionary literal expression.

msfterictraut · msfterictraut · commit eb1ef07671ba · 2022-08-09T09:05:51.000-07:00
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -12184,6 +12184,25 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
                 const unexpandedType = unexpandedTypeResult.type;
                 if (isAnyOrUnknown(unexpandedType)) {
                     addUnknown = false;
+                } else if (isClassInstance(unexpandedType) && ClassType.isTypedDictClass(unexpandedType)) {
+                    // Handle dictionary expansion for a TypedDict.
+                    if (strClassType && isInstantiableClass(strClassType)) {
+                        const strObject = ClassType.cloneAsInstance(strClassType);
+                        const tdEntries = getTypedDictMembersForClass(
+                            evaluatorInterface,
+                            unexpandedType,
+                            /* allowNarrowed */ true
+                        );
+
+                        tdEntries.forEach((entry, name) => {
+                            if (entry.isRequired || entry.isProvided) {
+                                keyTypes.push(ClassType.cloneWithLiteral(strObject, name));
+                                valueTypes.push(entry.valueType);
+                            }
+                        });
+
+                        addUnknown = false;
+                    }
                 } else {
                     // Verify that the type supports the `keys` and `__getitem__` methods.
                     // This protocol is defined in the _typeshed stub. If we can't find
diff --git a/packages/pyright-internal/src/tests/samples/typedDict21.py b/packages/pyright-internal/src/tests/samples/typedDict21.py
@@ -0,0 +1,33 @@
+# This sample tests the handling of dictionary expansion for TypedDicts.
+
+from typing_extensions import NotRequired, Required, TypedDict
+
+
+class TD1(TypedDict):
+    v1: Required[int]
+
+
+class TD2(TypedDict):
+    v2: Required[str]
+
+
+class TD3(TypedDict):
+    v1: NotRequired[int]
+
+
+class TD4(TD1, TD2):
+    ...
+
+
+td1: TD1 = {"v1": 0}
+td2: TD2 = {"v2": ""}
+td3_1: TD3 = {}
+td3_2: TD3 = {"v1": 0}
+
+td4_1: TD4 = {**td1, **td2}
+
+# This should generate an error because td3_1
+# does not include the required "v1" entry.
+td4_2: TD4 = {**td3_1, **td2}
+
+td4_3: TD4 = {**td3_2, **td2}
diff --git a/packages/pyright-internal/src/tests/typeEvaluator2.test.ts b/packages/pyright-internal/src/tests/typeEvaluator2.test.ts
@@ -1287,3 +1287,9 @@ test('TypedDict20', () => {
 
     TestUtils.validateResults(analysisResults, 0);
 });
+
+test('TypedDict21', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typedDict21.py']);
+
+    TestUtils.validateResults(analysisResults, 1);
+});
