From c285825deaa99320537e5ec5961f81b038f3127d Mon Sep 17 00:00:00 2001
From: Eric Traut <eric@traut.com>
Date: Tue, 6 Aug 2024 11:44:51 -0600
Subject: [PATCH] Allow `Final` and `ClassVar` to be combined in both
 directions within a dataclass. Previously, the `Final` qualifier needed to be
 the outermost. This addresses #8676. (#8678)

---
 .../src/analyzer/typeEvaluator.ts             | 24 ++++++++++++-------
 .../src/tests/samples/dataclass17.py          |  4 ++++
 .../src/tests/typeEvaluator4.test.ts          |  2 +-
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
index d96f5d09786c..b35a8dcc9553 100644
--- a/packages/pyright-internal/src/analyzer/typeEvaluator.ts
+++ b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -7607,24 +7607,30 @@ export function createTypeEvaluator(
         const typeArgs: TypeResultWithNode[] = [];
         let adjFlags = flags | EvalFlags.NoConvertSpecialForm;
 
-        if (options?.isFinalAnnotation) {
-            adjFlags |= EvalFlags.NoClassVar | EvalFlags.NoFinal;
-        } else if (options?.isClassVarAnnotation) {
-            adjFlags |= EvalFlags.NoClassVar;
-
+        const allowFinalClassVar = () => {
             // If the annotation is a variable within the body of a dataclass, a
-            // Final is allowed within the ClassVar annotation. In all other cases,
+            // Final is allowed with a ClassVar annotation. In all other cases,
             // it's disallowed.
-            let disallowFinal = true;
             const enclosingClassNode = ParseTreeUtils.getEnclosingClass(node, /* stopeAtFunction */ true);
             if (enclosingClassNode) {
                 const classTypeInfo = getTypeOfClass(enclosingClassNode);
                 if (classTypeInfo && ClassType.isDataClass(classTypeInfo.classType)) {
-                    disallowFinal = false;
+                    return true;
                 }
             }
+            return false;
+        };
+
+        if (options?.isFinalAnnotation) {
+            adjFlags |= EvalFlags.NoFinal;
+
+            if (!allowFinalClassVar()) {
+                adjFlags |= EvalFlags.NoClassVar;
+            }
+        } else if (options?.isClassVarAnnotation) {
+            adjFlags |= EvalFlags.NoClassVar;
 
-            if (disallowFinal) {
+            if (!allowFinalClassVar()) {
                 adjFlags |= EvalFlags.NoFinal;
             }
         } else {
diff --git a/packages/pyright-internal/src/tests/samples/dataclass17.py b/packages/pyright-internal/src/tests/samples/dataclass17.py
index 2a512ab4ebaf..81d64b73e374 100644
--- a/packages/pyright-internal/src/tests/samples/dataclass17.py
+++ b/packages/pyright-internal/src/tests/samples/dataclass17.py
@@ -11,6 +11,7 @@ class A:
     b: Final[str] = ""
     c: ClassVar[Final[int]] = 0
     d: ClassVar[Final] = 0
+    e: Final[ClassVar[int]] = 0
 
 
 a = A(1)
@@ -29,3 +30,6 @@ class A:
 
 # This should generate an error.
 A.d = 0
+
+# This should generate an error.
+A.e = 0
diff --git a/packages/pyright-internal/src/tests/typeEvaluator4.test.ts b/packages/pyright-internal/src/tests/typeEvaluator4.test.ts
index baa923db4131..2119c653fcd3 100644
--- a/packages/pyright-internal/src/tests/typeEvaluator4.test.ts
+++ b/packages/pyright-internal/src/tests/typeEvaluator4.test.ts
@@ -363,7 +363,7 @@ test('DataClass16', () => {
 test('DataClass17', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclass17.py']);
 
-    TestUtils.validateResults(analysisResults, 5);
+    TestUtils.validateResults(analysisResults, 6);
 });
 
 test('DataClassReplace1', () => {