Added experimental support for draft PEP 764: Inlined typed dictionar…

…ies. (#9350)

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

@@ -72,6 +72,9 @@ export interface DeclarationBase {
     // The declaration is within an except clause of a try
     // statement. We may want to ignore such declarations.
     isInExceptSuite: boolean;
+
+    // This declaration is within an inlined TypedDict definition.
+    isInInlinedTypedDict?: boolean;
 }
 
 export interface IntrinsicDeclaration extends DeclarationBase {

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

@@ -174,6 +174,7 @@ import {
     assignToTypedDict,
     assignTypedDictToTypedDict,
     createTypedDictType,
+    createTypedDictTypeInlined,
     getTypedDictDictEquivalent,
     getTypedDictMappingEquivalent,
     getTypedDictMembersForClass,
@@ -368,6 +369,7 @@ interface GetTypeArgsOptions {
     hasCustomClassGetItem?: boolean;
     isFinalAnnotation?: boolean;
     isClassVarAnnotation?: boolean;
+    supportsTypedDictTypeArg?: boolean;
 }
 
 interface MatchArgsToParamsResult {
@@ -654,6 +656,7 @@ export function createTypeEvaluator(
     let strClass: Type | undefined;
     let dictClass: Type | undefined;
     let moduleTypeClass: Type | undefined;
+    let typedDictClass: Type | undefined;
     let typedDictPrivateClass: Type | undefined;
     let supportsKeysAndGetItemClass: Type | undefined;
     let mappingClass: Type | undefined;
@@ -1022,6 +1025,7 @@ export function createTypeEvaluator(
             strClass = getBuiltInType(node, 'str');
             dictClass = getBuiltInType(node, 'dict');
             moduleTypeClass = getTypingType(node, 'ModuleType');
+            typedDictClass = getTypingType(node, 'TypedDict');
             typedDictPrivateClass = getTypingType(node, '_TypedDict');
             awaitableClass = getTypingType(node, 'Awaitable');
             mappingClass = getTypingType(node, 'Mapping');
@@ -7523,11 +7527,17 @@ export function createTypeEvaluator(
                     const isClassVarAnnotation =
                         isInstantiableClass(concreteSubtype) && ClassType.isBuiltIn(concreteSubtype, 'ClassVar');
 
+                    // This feature is currently experimental.
+                    const supportsTypedDictTypeArg =
+                        AnalyzerNodeInfo.getFileInfo(node).diagnosticRuleSet.enableExperimentalFeatures &&
+                        ClassType.isBuiltIn(concreteSubtype, 'TypedDict');
+
                     let typeArgs = getTypeArgs(node, flags, {
                         isAnnotatedClass,
                         hasCustomClassGetItem: hasCustomClassGetItem || !isGenericClass,
                         isFinalAnnotation,
                         isClassVarAnnotation,
+                        supportsTypedDictTypeArg,
                     });
 
                     if (!isAnnotatedClass) {
@@ -8014,7 +8024,7 @@ export function createTypeEvaluator(
                     node: expr,
                 };
             } else {
-                typeResult = getTypeArg(expr, adjFlags);
+                typeResult = getTypeArg(expr, adjFlags, !!options?.supportsTypedDictTypeArg && argIndex === 0);
             }
 
             return typeResult;
@@ -8120,7 +8130,7 @@ export function createTypeEvaluator(
         return undefined;
     }
 
-    function getTypeArg(node: ExpressionNode, flags: EvalFlags): TypeResultWithNode {
+    function getTypeArg(node: ExpressionNode, flags: EvalFlags, supportsDictExpression: boolean): TypeResultWithNode {
         let typeResult: TypeResultWithNode;
 
         let adjustedFlags =
@@ -8142,6 +8152,18 @@ export function createTypeEvaluator(
 
             // Set the node's type so it isn't reevaluated later.
             setTypeResultForNode(node, { type: UnknownType.create() });
+        } else if (node.nodeType === ParseNodeType.Dictionary && supportsDictExpression) {
+            const inlinedTypeDict =
+                typedDictClass && isInstantiableClass(typedDictClass)
+                    ? createTypedDictTypeInlined(evaluatorInterface, node, typedDictClass)
+                    : undefined;
+            const keyTypeFallback = strClass && isInstantiableClass(strClass) ? strClass : UnknownType.create();
+
+            typeResult = {
+                type: keyTypeFallback,
+                inlinedTypeDict,
+                node,
+            };
         } else {
             typeResult = { ...getTypeOfExpression(node, adjustedFlags), node };
 
@@ -20577,11 +20599,17 @@ export function createTypeEvaluator(
 
                 case 'TypedDict': {
                     if ((flags & (EvalFlags.NoNonTypeSpecialForms | EvalFlags.TypeExpression)) !== 0) {
-                        addDiagnostic(
-                            DiagnosticRule.reportInvalidTypeForm,
-                            LocMessage.typedDictNotAllowed(),
-                            errorNode
-                        );
+                        const isInlinedTypedDict =
+                            AnalyzerNodeInfo.getFileInfo(errorNode).diagnosticRuleSet.enableExperimentalFeatures &&
+                            !!typeArgs;
+
+                        if (!isInlinedTypedDict) {
+                            addDiagnostic(
+                                DiagnosticRule.reportInvalidTypeForm,
+                                LocMessage.typedDictNotAllowed(),
+                                errorNode
+                            );
+                        }
                     }
                     isValidTypeForm = false;
                     break;
@@ -20746,7 +20774,22 @@ export function createTypeEvaluator(
                 minTypeArgCount = firstDefaultParamIndex;
             }
 
-            if (typeArgCount > typeParams.length) {
+            // Classes that accept inlined type dict type args allow only one.
+            if (typeArgs.length > 0 && typeArgs[0].inlinedTypeDict) {
+                if (typeArgs.length > 1) {
+                    addDiagnostic(
+                        DiagnosticRule.reportInvalidTypeArguments,
+                        LocMessage.typeArgsTooMany().format({
+                            name: classType.priv.aliasName || classType.shared.name,
+                            expected: 1,
+                            received: typeArgCount,
+                        }),
+                        typeArgs[1].node
+                    );
+                }
+
+                return { type: typeArgs[0].inlinedTypeDict };
+            } else if (typeArgCount > typeParams.length) {
                 if (!ClassType.isPartiallyEvaluated(classType) && !ClassType.isTupleClass(classType)) {
                     if (typeParams.length === 0) {
                         isValidTypeForm = false;
@@ -21862,12 +21905,18 @@ export function createTypeEvaluator(
                             declaration.node.parent?.nodeType === ParseNodeType.MemberAccess
                                 ? declaration.node.parent
                                 : declaration.node;
+                        const allowClassVar = ParseTreeUtils.isClassVarAllowedForAssignmentTarget(declNode);
+                        const allowFinal = ParseTreeUtils.isFinalAllowedForAssignmentTarget(declNode);
+                        const allowRequired =
+                            ParseTreeUtils.isRequiredAllowedForAssignmentTarget(declNode) ||
+                            !!declaration.isInInlinedTypedDict;
+
                         declaredType = getTypeOfAnnotation(typeAnnotationNode, {
                             varTypeAnnotation: true,
-                            allowClassVar: ParseTreeUtils.isClassVarAllowedForAssignmentTarget(declNode),
-                            allowFinal: ParseTreeUtils.isFinalAllowedForAssignmentTarget(declNode),
-                            allowRequired: ParseTreeUtils.isRequiredAllowedForAssignmentTarget(declNode),
-                            allowReadOnly: ParseTreeUtils.isRequiredAllowedForAssignmentTarget(declNode),
+                            allowClassVar,
+                            allowFinal,
+                            allowRequired,
+                            allowReadOnly: allowRequired,
                             enforceClassTypeVarScope: declaration.isDefinedByMemberAccess,
                         });
                     }

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

@@ -221,6 +221,9 @@ export interface TypeResult<T extends Type = Type> {
     // Type consistency errors detected when evaluating this type.
     typeErrors?: boolean | undefined;
 
+    // For inlined TypedDict definitions.
+    inlinedTypeDict?: ClassType;
+
     // Used for getTypeOfBoundMember to indicate that class
     // that declares the member.
     classType?: ClassType | UnknownType | AnyType;

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

@@ -259,6 +259,34 @@ export function createTypedDictType(
     return classType;
 }
 
+// Creates a new anonymous TypedDict class from an inlined dict[{}] type annotation.
+export function createTypedDictTypeInlined(
+    evaluator: TypeEvaluator,
+    dictNode: DictionaryNode,
+    typedDictClass: ClassType
+): ClassType {
+    const fileInfo = AnalyzerNodeInfo.getFileInfo(dictNode);
+    const className = '<TypedDict>';
+
+    const classType = ClassType.createInstantiable(
+        className,
+        ParseTreeUtils.getClassFullName(dictNode, fileInfo.moduleName, className),
+        fileInfo.moduleName,
+        fileInfo.fileUri,
+        ClassTypeFlags.TypedDictClass,
+        ParseTreeUtils.getTypeSourceId(dictNode),
+        /* declaredMetaclass */ undefined,
+        typedDictClass.shared.effectiveMetaclass
+    );
+    classType.shared.baseClasses.push(typedDictClass);
+    computeMroLinearization(classType);
+
+    getTypedDictFieldsFromDictSyntax(evaluator, dictNode, ClassType.getSymbolTable(classType), /* isInline */ true);
+    synthesizeTypedDictClassMethods(evaluator, dictNode, classType);
+
+    return classType;
+}
+
 export function synthesizeTypedDictClassMethods(
     evaluator: TypeEvaluator,
     node: ClassNode | ExpressionNode,
@@ -964,6 +992,7 @@ function getTypedDictFieldsFromDictSyntax(
             range: convertOffsetsToRange(entry.d.keyExpr.start, TextRange.getEnd(entry.d.keyExpr), fileInfo.lines),
             moduleName: fileInfo.moduleName,
             isInExceptSuite: false,
+            isInInlinedTypedDict: true,
         };
         newSymbol.addDeclaration(declaration);
 

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

@@ -0,0 +1,40 @@
+# This sample tests support for inlined TypedDict definitions.
+
+from typing import NotRequired, ReadOnly, Required, TypedDict
+
+
+td1: TypedDict[{"a": int, "b": str}] = {"a": 0, "b": ""}
+
+td2: TypedDict[{"a": TypedDict[{"b": int}]}] = {"a": {"b": 0}}
+
+td3: TypedDict[{"a": "list[float]"}] = {"a": [3]}
+
+td4: TypedDict[
+    {"a": NotRequired[int], "b": Required[int], "c": NotRequired[ReadOnly[int]]}
+] = {"b": 3}
+
+# This should generate an error because dictionary comprehensions
+# are not allowed.
+err1: TypedDict[{"a": int for _ in range(1)}]
+
+# This should generate an error because unpacked dictionary
+# entries are not allowed.
+err2: TypedDict[{**{"a": int}}]
+
+# This should generate an error because an extra type argument is provided.
+err3: TypedDict[{"a": int}, str]
+
+# This should generate an error because TypedDict cannot be used without
+# a subscript in this context.
+err4: TypedDict
+
+# This should generate an error because a dict expression is not a
+# valid type expression by itself.
+err5: TypedDict[{"a": {"b": int}}] = {"a": {"b": 0}}
+
+
+def func1(val: TypedDict[{"a": int}]) -> TypedDict[{"a": int}]:
+    return {"a": val["a"] + 1}
+
+
+func1({"a": 3})

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

@@ -767,6 +767,14 @@ test('TypedDict24', () => {
     TestUtils.validateResults(analysisResults, 1);
 });
 
+test('TypedDictInline1', () => {
+    const configOptions = new ConfigOptions(Uri.empty());
+    configOptions.diagnosticRuleSet.enableExperimentalFeatures = true;
+
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typedDictInline1.py'], configOptions);
+    TestUtils.validateResults(analysisResults, 6);
+});
+
 test('ClassVar1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['classVar1.py']);