Implemented support for PEP 655: Marking individual TypedDict items as required or potentially-missing.

Author: msfterictraut

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

@@ -1426,7 +1426,7 @@ export class Binder extends ParseTreeWalker {
     }
 
     visitImportFrom(node: ImportFromNode): boolean {
-        const typingSymbolsOfInterest = ['Final', 'TypeAlias', 'ClassVar'];
+        const typingSymbolsOfInterest = ['Final', 'TypeAlias', 'ClassVar', 'Required', 'NotRequired'];
         const importInfo = AnalyzerNodeInfo.getImportInfo(node.module);
 
         let resolvedPath = '';
@@ -2857,6 +2857,8 @@ export class Binder extends ParseTreeWalker {
                         node: target,
                         isConstant: isConstantName(name.value),
                         isFinal: finalInfo.isFinal,
+                        isRequired: this._isRequiredAnnotation(typeAnnotationNode),
+                        isNotRequired: this._isNotRequiredAnnotation(typeAnnotationNode),
                         typeAliasAnnotation: isExplicitTypeAlias ? typeAnnotation : undefined,
                         typeAliasName: isExplicitTypeAlias ? target : undefined,
                         path: this._fileInfo.filePath,
@@ -2965,7 +2967,7 @@ export class Binder extends ParseTreeWalker {
     }
 
     // Determines if the specified type annotation expression is a "Final".
-    // It returns two boolean values indicating if the expression is a "Final"
+    // It returns a value indicating whether the expression is a "Final"
     // expression and whether it's a "raw" Final with no type arguments.
     private _isAnnotationFinal(typeAnnotation: ExpressionNode | undefined): FinalInfo {
         let isFinal = false;
@@ -2992,6 +2994,28 @@ export class Binder extends ParseTreeWalker {
         return { isFinal, finalTypeNode };
     }
 
+    // Determines if the specified type annotation is wrapped in a "Required".
+    private _isRequiredAnnotation(typeAnnotation: ExpressionNode | undefined): boolean {
+        if (typeAnnotation && typeAnnotation.nodeType === ParseNodeType.Index && typeAnnotation.items.length === 1) {
+            if (this._isTypingAnnotation(typeAnnotation.baseExpression, 'Required')) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    // Determines if the specified type annotation is wrapped in a "NotRequired".
+    private _isNotRequiredAnnotation(typeAnnotation: ExpressionNode | undefined): boolean {
+        if (typeAnnotation && typeAnnotation.nodeType === ParseNodeType.Index && typeAnnotation.items.length === 1) {
+            if (this._isTypingAnnotation(typeAnnotation.baseExpression, 'NotRequired')) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     private _isAnnotationTypeAlias(typeAnnotation: ExpressionNode | undefined) {
         if (!typeAnnotation) {
             return false;

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

@@ -112,6 +112,12 @@ export interface VariableDeclaration extends DeclarationBase {
     // constant in that reassignment is not permitted)?
     isFinal?: boolean;
 
+    // Is the declaration annotated with "Required"?
+    isRequired?: boolean;
+
+    // Is the declaration annotated with "NotRequired"?
+    isNotRequired?: boolean;
+
     // Points to the "TypeAlias" annotation described in PEP 613.
     typeAliasAnnotation?: ExpressionNode;
 

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

@@ -44,7 +44,7 @@ export const enum SymbolFlags {
     // set when accessed through a class instance.
     ClassVar = 1 << 7,
 
-    // // Indicates that the symbol is in __all__.
+    // Indicates that the symbol is in __all__.
     InDunderAll = 1 << 8,
 }
 

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

@@ -39,3 +39,11 @@ export function isTypedDictMemberAccessedThroughIndex(symbol: Symbol): boolean {
 export function isFinalVariable(symbol: Symbol): boolean {
     return symbol.getDeclarations().some((decl) => isFinalVariableDeclaration(decl));
 }
+
+export function isRequiredTypedDictVariable(symbol: Symbol) {
+    return symbol.getDeclarations().some((decl) => decl.type === DeclarationType.Variable && !!decl.isRequired);
+}
+
+export function isNotRequiredTypedDictVariable(symbol: Symbol) {
+    return symbol.getDeclarations().some((decl) => decl.type === DeclarationType.Variable && !!decl.isNotRequired);
+}

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

@@ -115,7 +115,7 @@ import * as ScopeUtils from './scopeUtils';
 import { evaluateStaticBoolExpression } from './staticExpressions';
 import { indeterminateSymbolId, Symbol, SymbolFlags } from './symbol';
 import { isConstantName, isDunderName, isPrivateOrProtectedName } from './symbolNameUtils';
-import { getLastTypedDeclaredForSymbol, isFinalVariable } from './symbolUtils';
+import { getLastTypedDeclaredForSymbol, isFinalVariable, isNotRequiredTypedDictVariable, isRequiredTypedDictVariable } from './symbolUtils';
 import { PrintableType, TracePrinter } from './tracePrinter';
 import {
     CachedType,
@@ -10220,6 +10220,51 @@ export function createTypeEvaluator(
         return ClassType.cloneForSpecialization(classType, [convertToInstance(typeArg)], !!typeArgs);
     }
 
+    function createRequiredType(
+        classType: ClassType,
+        errorNode: ParseNode,
+        isRequired: boolean,
+        typeArgs: TypeResult[] | undefined
+    ): Type {
+        if (!typeArgs || typeArgs.length !== 1) {
+            addError(
+                isRequired ? Localizer.Diagnostic.requiredArgCount() : Localizer.Diagnostic.notRequiredArgCount(),
+                errorNode
+            );
+            return classType;
+        }
+
+        const typeArgType = typeArgs[0].type;
+
+        // Make sure this is used only in a dataclass.
+        const containingClassNode = ParseTreeUtils.getEnclosingClass(errorNode, /* stopAtFunction */ true);
+        const classTypeInfo = containingClassNode ? getTypeOfClass(containingClassNode) : undefined;
+
+        let isUsageLegal = false;
+
+        if (classTypeInfo && isClass(classTypeInfo.classType) && ClassType.isTypedDictClass(classTypeInfo.classType)) {
+            // The only legal usage is when used in a type annotation statement.
+            if (
+                errorNode.parent?.nodeType === ParseNodeType.TypeAnnotation &&
+                errorNode.parent.typeAnnotation === errorNode
+            ) {
+                isUsageLegal = true;
+            }
+        }
+
+        if (!isUsageLegal) {
+            addError(
+                isRequired
+                    ? Localizer.Diagnostic.requiredNotInTypedDict()
+                    : Localizer.Diagnostic.notRequiredNotInTypedDict(),
+                errorNode
+            );
+            return ClassType.cloneForSpecialization(classType, [convertToInstance(typeArgType)], !!typeArgs);
+        }
+
+        return typeArgType;
+    }
+
     function createUnpackType(errorNode: ParseNode, typeArgs: TypeResult[] | undefined): Type {
         if (!typeArgs || typeArgs.length !== 1) {
             addError(Localizer.Diagnostic.unpackArgCount(), errorNode);
@@ -10435,7 +10480,7 @@ export function createTypeEvaluator(
     function transformTypeForPossibleEnumClass(node: NameNode, typeOfExpr: Type): Type {
         // If the node is within a class that derives from the metaclass
         // "EnumMeta", we need to treat assignments differently.
-        const enclosingClassNode = ParseTreeUtils.getEnclosingClass(node, true);
+        const enclosingClassNode = ParseTreeUtils.getEnclosingClass(node, /* stopAtFunction */ true);
         if (enclosingClassNode) {
             const enumClassInfo = getTypeOfClass(enclosingClassNode);
 
@@ -10630,6 +10675,8 @@ export function createTypeEvaluator(
             Concatenate: { alias: '', module: 'builtins' },
             TypeGuard: { alias: '', module: 'builtins' },
             Unpack: { alias: '', module: 'builtins' },
+            Required: { alias: '', module: 'builtins' },
+            NotRequired: { alias: '', module: 'builtins' },
         };
 
         const aliasMapEntry = specialTypes[assignedName];
@@ -11462,7 +11509,7 @@ export function createTypeEvaluator(
 
         // There was no cached type, so create a new one.
         // Retrieve the containing class node if the function is a method.
-        const containingClassNode = ParseTreeUtils.getEnclosingClass(node, true);
+        const containingClassNode = ParseTreeUtils.getEnclosingClass(node, /* stopAtFunction */ true);
         let containingClassType: ClassType | undefined;
         if (containingClassNode) {
             const classInfo = getTypeOfClass(containingClassNode);
@@ -15808,6 +15855,11 @@ export function createTypeEvaluator(
                 case 'Unpack': {
                     return createUnpackType(errorNode, typeArgs);
                 }
+
+                case 'Required':
+                case 'NotRequired': {
+                    return createRequiredType(classType, errorNode, aliasedName === 'Required', typeArgs);
+                }
             }
         }
 
@@ -19924,9 +19976,18 @@ export function createTypeEvaluator(
                 // Only variables (not functions, classes, etc.) are considered.
                 const lastDecl = getLastTypedDeclaredForSymbol(symbol);
                 if (lastDecl && lastDecl.type === DeclarationType.Variable) {
+                    const valueType = getDeclaredTypeOfSymbol(symbol) || UnknownType.create();
+                    let isRequired = !ClassType.isCanOmitDictValues(classType);
+
+                    if (isRequiredTypedDictVariable(symbol)) {
+                        isRequired = true;
+                    } else if (isNotRequiredTypedDictVariable(symbol)) {
+                        isRequired = false;
+                    }
+
                     keyMap.set(name, {
-                        valueType: getDeclaredTypeOfSymbol(symbol) || UnknownType.create(),
-                        isRequired: !ClassType.isCanOmitDictValues(classType),
+                        valueType,
+                        isRequired,
                         isProvided: false,
                     });
                 }

packages/pyright-internal/src/localization/localize.ts

@@ -451,6 +451,8 @@ export namespace Localizer {
         export const noReturnContainsReturn = () => getRawString('Diagnostic.noReturnContainsReturn');
         export const noReturnContainsYield = () => getRawString('Diagnostic.noReturnContainsYield');
         export const noReturnReturnsNone = () => getRawString('Diagnostic.noReturnReturnsNone');
+        export const notRequiredArgCount = () => getRawString('Diagnostic.notRequiredArgCount');
+        export const notRequiredNotInTypedDict = () => getRawString('Diagnostic.notRequiredNotInTypedDict');
         export const objectNotCallable = () =>
             new ParameterizedString<{ type: string }>(getRawString('Diagnostic.objectNotCallable'));
         export const obscuredClassDeclaration = () =>
@@ -535,6 +537,8 @@ export namespace Localizer {
         export const recursiveDefinition = () =>
             new ParameterizedString<{ name: string }>(getRawString('Diagnostic.recursiveDefinition'));
         export const relativeImportNotAllowed = () => getRawString('Diagnostic.relativeImportNotAllowed');
+        export const requiredArgCount = () => getRawString('Diagnostic.requiredArgCount');
+        export const requiredNotInTypedDict = () => getRawString('Diagnostic.requiredNotInTypedDict');
         export const returnMissing = () =>
             new ParameterizedString<{ returnType: string }>(getRawString('Diagnostic.returnMissing'));
         export const returnOutsideFunction = () => getRawString('Diagnostic.returnOutsideFunction');

packages/pyright-internal/src/localization/package.nls.en-us.json

@@ -218,6 +218,8 @@
         "noReturnContainsReturn": "Function with declared return type \"NoReturn\" cannot include a return statement",
         "noReturnContainsYield": "Function with declared return type \"NoReturn\" cannot include a yield statement",
         "noReturnReturnsNone": "Function with declared type of \"NoReturn\" cannot return \"None\"",
+        "notRequiredArgCount": "Expected a single type argument after \"NotRequired\"",
+        "notRequiredNotInTypedDict": "\"NotRequired\" is allowed only within TypedDict",
         "objectNotCallable": "Object of type \"{type}\" is not callable",
         "obscuredClassDeclaration": "Class declaration \"{name}\" is obscured by a declaration of the same name",
         "obscuredFunctionDeclaration": "Function declaration \"{name}\" is obscured by a declaration of the same name",
@@ -264,6 +266,8 @@
         "raiseParams": "\"raise\" requires one or more parameters when used outside of except clause",
         "relativeImportNotAllowed": "Relative imports cannot be used with \"import .a\" form; use \"from . import a\" instead",
         "recursiveDefinition": "Type of \"{name}\" could not be determined because it refers to itself",
+        "requiredArgCount": "Expected a single type argument after \"Required\"",
+        "requiredNotInTypedDict": "\"Required\" is allowed only within TypedDict",
         "returnOutsideFunction": "\"return\" can be used only within a function",
         "returnMissing": "Function with declared type of \"{returnType}\" must return value",
         "returnTypeContravariant": "Contravariant type variable cannot be used in return type",

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

@@ -0,0 +1,44 @@
+# This sample tests the handling of Required and NotRequired
+# (PEP 655) in TypedDict definitions.
+
+from typing import Annotated, NotRequired, Required, TypedDict
+
+
+class TD1(TypedDict):
+    a: Required[int]
+    b: NotRequired[int]
+
+    # This should generate an error because NotRequired can't be
+    # used in this context.
+    c: NotRequired[NotRequired[int]]
+
+    # This should generate an error because Required can't be
+    # used in this context.
+    d: Required[Required[int]]
+
+    e: NotRequired[Annotated[int, "hi"]]
+
+    # This should generate an error because it's missing type args.
+    f: Required
+
+    # This should generate an error because it's missing type args.
+    g: NotRequired
+
+
+# This should generate an error because Required can't be
+# used in this context.
+x: Required[int]
+
+# This should generate an error because NotRequired can't be
+# used in this context.
+y: Required[int]
+
+
+class Foo:
+    # This should generate an error because Required can't be
+    # used in this context.
+    x: Required[int]
+
+    # This should generate an error because NotRequired can't be
+    # used in this context.
+    y: Required[int]

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

@@ -0,0 +1,53 @@
+# This sample tests the handling of Required and NotRequired
+# (PEP 655) in TypedDict definitions.
+
+from typing import Literal, NotRequired, Optional, Required, Type, TypedDict
+
+
+class TD1(TypedDict, total=False):
+    a: Required[int]
+    b: NotRequired[str]
+    c: Required[int | str]
+    d: Required[Optional[str]]
+    e: Required[Literal[1, 2, 3]]
+    f: Required[None]
+    g: Required[Type[int]]
+
+
+td1_1: TD1 = {"a": 3, "c": "hi", "d": None, "e": 3, "f": None, "g": int}
+
+# This should generate an error because a is missing.
+td1_2: TD1 = {"c": "hi", "d": None, "e": 3, "f": None, "g": int}
+
+# This should generate an error because c is missing.
+td1_3: TD1 = {"a": 3, "d": None, "e": 3, "f": None, "g": int}
+
+# This should generate an error because d is missing.
+td1_4: TD1 = {"a": 3, "c": "hi", "e": 3, "f": None, "g": int}
+
+# This should generate an error because e is missing.
+td1_5: TD1 = {"a": 3, "c": "hi", "d": None, "f": None, "g": int}
+
+# This should generate an error because f is missing.
+td1_6: TD1 = {"a": 3, "c": "hi", "d": None, "e": 3, "g": int}
+
+# This should generate an error because g is missing.
+td1_7: TD1 = {"a": 3, "c": "hi", "d": None, "e": 3, "f": None}
+
+
+class TD2(TypedDict, total=True):
+    a: Required[int]
+    b: NotRequired[str]
+    c: Required[int | str]
+    d: NotRequired[Optional[str]]
+    e: NotRequired[Literal[1, 2, 3]]
+    f: NotRequired[None]
+    g: NotRequired[Type[int]]
+
+
+td2_1: TD2 = {"a": 3, "c": "hi", "d": None, "e": 3, "f": None, "g": int}
+
+td2_2: TD2 = {"a": 3, "c": "hi"}
+
+# This should generate an error because c is missing.
+td2_3: TD2 = {"a": 3}

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

@@ -705,6 +705,24 @@ test('TypedDict12', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Required1', () => {
+    // Analyze with Python 3.10 settings.
+    const configOptions = new ConfigOptions('.');
+    configOptions.defaultPythonVersion = PythonVersion.V3_10;
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['required1.py'], configOptions);
+
+    TestUtils.validateResults(analysisResults, 8);
+});
+
+test('Required2', () => {
+    // Analyze with Python 3.10 settings.
+    const configOptions = new ConfigOptions('.');
+    configOptions.defaultPythonVersion = PythonVersion.V3_10;
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['required2.py'], configOptions);
+
+    TestUtils.validateResults(analysisResults, 7);
+});
+
 test('Metaclass1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['metaclass1.py']);
     TestUtils.validateResults(analysisResults, 0);

packages/pyright-internal/typeshed-fallback/stdlib/typing.pyi

@@ -62,6 +62,8 @@ if sys.version_info >= (3, 10):
     Concatenate: _SpecialForm = ...
     TypeAlias: _SpecialForm = ...
     TypeGuard: _SpecialForm = ...
+    Required: _SpecialForm = ...
+    NotRequired: _SpecialForm = ...
 
     class TypeVarTuple:
         __name__: str

packages/pyright-internal/typeshed-fallback/stubs/typing-extensions/typing_extensions.pyi

@@ -122,3 +122,7 @@ class TypeVarTuple:
 
 # PEP 647
 TypeGuard: _SpecialForm = ...
+
+# PEP 655
+Required: _SpecialForm = ...
+NotRequired: _SpecialForm = ...