From 07c2b9ed371c43f39a636ec70a8c9bc7248c6435 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 29 Oct 2024 08:38:07 -0700 Subject: [PATCH] Added experimental support for draft PEP 764: Inlined typed dictionaries. --- .../src/analyzer/declaration.ts | 3 + .../src/analyzer/typeEvaluator.ts | 73 ++++++++++++++++--- .../src/analyzer/typeEvaluatorTypes.ts | 3 + .../src/analyzer/typedDicts.ts | 29 ++++++++ .../src/tests/samples/typedDictInline1.py | 40 ++++++++++ .../src/tests/typeEvaluator7.test.ts | 8 ++ 6 files changed, 144 insertions(+), 12 deletions(-) create mode 100644 packages/pyright-internal/src/tests/samples/typedDictInline1.py diff --git a/packages/pyright-internal/src/analyzer/declaration.ts b/packages/pyright-internal/src/analyzer/declaration.ts index 1eeffc929d4b..dca2d6456e03 100644 --- a/packages/pyright-internal/src/analyzer/declaration.ts +++ b/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 { diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts index 78b9d8f2aba2..b41f7478fc16 100644 --- a/packages/pyright-internal/src/analyzer/typeEvaluator.ts +++ b/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, }); } diff --git a/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts b/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts index 0d57d180ee15..572b38d7b7ec 100644 --- a/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts +++ b/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts @@ -221,6 +221,9 @@ export interface TypeResult { // 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; diff --git a/packages/pyright-internal/src/analyzer/typedDicts.ts b/packages/pyright-internal/src/analyzer/typedDicts.ts index a1fdce373fff..caa2fd9b1336 100644 --- a/packages/pyright-internal/src/analyzer/typedDicts.ts +++ b/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 = ''; + + 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); diff --git a/packages/pyright-internal/src/tests/samples/typedDictInline1.py b/packages/pyright-internal/src/tests/samples/typedDictInline1.py new file mode 100644 index 000000000000..8e1be9da0610 --- /dev/null +++ b/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}) diff --git a/packages/pyright-internal/src/tests/typeEvaluator7.test.ts b/packages/pyright-internal/src/tests/typeEvaluator7.test.ts index 08a2a7ce04f0..0ad4003e90d8 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator7.test.ts +++ b/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']);