Added check for invalid use of ClassVar qualifier within a NamedTuple or TypedDict attribute annotation. This addresses #9526. (#9540)

Author: erictraut

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

@@ -1178,15 +1178,6 @@ export function isFinalAllowedForAssignmentTarget(targetNode: ExpressionNode): b
     return false;
 }
 
-export function isClassVarAllowedForAssignmentTarget(targetNode: ExpressionNode): boolean {
-    const classNode = getEnclosingClass(targetNode, /* stopAtFunction */ true);
-    if (!classNode) {
-        return false;
-    }
-
-    return true;
-}
-
 export function isRequiredAllowedForAssignmentTarget(targetNode: ExpressionNode): boolean {
     const classNode = getEnclosingClass(targetNode, /* stopAtFunction */ true);
     if (!classNode) {

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

@@ -4360,7 +4360,7 @@ export function createTypeEvaluator(
                 let annotationType: Type | undefined = getTypeOfAnnotation(target.d.annotation, {
                     varTypeAnnotation: true,
                     allowFinal: ParseTreeUtils.isFinalAllowedForAssignmentTarget(target.d.valueExpr),
-                    allowClassVar: ParseTreeUtils.isClassVarAllowedForAssignmentTarget(target.d.valueExpr),
+                    allowClassVar: isClassVarAllowedForAssignmentTarget(target.d.valueExpr),
                 });
 
                 if (annotationType) {
@@ -4433,6 +4433,21 @@ export function createTypeEvaluator(
         }
     }
 
+    function isClassVarAllowedForAssignmentTarget(targetNode: ExpressionNode): boolean {
+        const classNode = ParseTreeUtils.getEnclosingClass(targetNode, /* stopAtFunction */ true);
+        if (!classNode) {
+            return false;
+        }
+
+        // ClassVar is not allowed in a TypedDict or a NamedTuple class.
+        const classType = getTypeOfClass(classNode)?.classType;
+        if (!classType) {
+            return false;
+        }
+
+        return !ClassType.isTypedDictClass(classType) && !classType.shared.namedTupleEntries;
+    }
+
     function verifyRaiseExceptionType(node: ExpressionNode, allowNone: boolean) {
         const baseExceptionType = getBuiltInType(node, 'BaseException');
         const exceptionType = getTypeOfExpression(node).type;
@@ -19924,7 +19939,7 @@ export function createTypeEvaluator(
             const annotationType = getTypeOfAnnotation(node.d.annotation, {
                 varTypeAnnotation: true,
                 allowFinal: ParseTreeUtils.isFinalAllowedForAssignmentTarget(node.d.valueExpr),
-                allowClassVar: ParseTreeUtils.isClassVarAllowedForAssignmentTarget(node.d.valueExpr),
+                allowClassVar: isClassVarAllowedForAssignmentTarget(node.d.valueExpr),
             });
 
             writeTypeCache(node.d.valueExpr, { type: annotationType }, EvalFlags.None);
@@ -20073,7 +20088,7 @@ export function createTypeEvaluator(
                     getTypeOfAnnotation(annotationNode, {
                         varTypeAnnotation: true,
                         allowFinal: ParseTreeUtils.isFinalAllowedForAssignmentTarget(annotationParent.d.leftExpr),
-                        allowClassVar: ParseTreeUtils.isClassVarAllowedForAssignmentTarget(annotationParent.d.leftExpr),
+                        allowClassVar: isClassVarAllowedForAssignmentTarget(annotationParent.d.leftExpr),
                     });
                 } else {
                     evaluateTypesForAssignmentStatement(annotationParent);
@@ -22048,7 +22063,7 @@ export function createTypeEvaluator(
                             declaration.node.parent?.nodeType === ParseNodeType.MemberAccess
                                 ? declaration.node.parent
                                 : declaration.node;
-                        const allowClassVar = ParseTreeUtils.isClassVarAllowedForAssignmentTarget(declNode);
+                        const allowClassVar = isClassVarAllowedForAssignmentTarget(declNode);
                         const allowFinal = ParseTreeUtils.isFinalAllowedForAssignmentTarget(declNode);
                         const allowRequired =
                             ParseTreeUtils.isRequiredAllowedForAssignmentTarget(declNode) ||

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

@@ -0,0 +1,18 @@
+# This sample tests that a ClassVar is disallowed when used in a 
+# NamedTuple or TypedDict class as reflected in the runtime.
+
+from typing import ClassVar, NamedTuple, TypedDict
+
+class NT1(NamedTuple):
+    # This should generate an error.
+    x: ClassVar
+
+    # This should generate an error.
+    y: ClassVar[int]
+
+class TD1(TypedDict):
+    # This should generate an error.
+    x: ClassVar
+
+    # This should generate an error.
+    y: ClassVar[int]

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

@@ -15,9 +15,6 @@ class DataTuple(NamedTuple):
     def _m(self):
         pass
 
-    # ClassVar variables should not be included.
-    class_var: ClassVar[int] = 4
-
     id: int
     aid: Other
     value: str = ""

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

@@ -805,6 +805,12 @@ test('ClassVar5', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('ClassVar6', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['classVar6.py']);
+
+    TestUtils.validateResults(analysisResults, 4);
+});
+
 test('TypeVar1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeVar1.py']);