From 0b3073f7a045ab12962b59e2db41379f7cd0eeab Mon Sep 17 00:00:00 2001
From: Eric Traut <eric@traut.com>
Date: Mon, 2 Dec 2024 06:00:00 -0800
Subject: [PATCH] Added check for importing a `Final` variable from another
 module and then attempting to overwrite it. This partially addresses
 https://github.com/microsoft/pylance-release/issues/6455.

---
 .../pyright-internal/src/analyzer/checker.ts  | 32 ++++++++++++++++++-
 .../src/analyzer/typeEvaluatorTypes.ts        |  4 +--
 .../src/tests/samples/final7.py               |  8 +++++
 .../src/tests/samples/final8.py               | 20 ++++++++++++
 .../src/tests/typeEvaluator4.test.ts          |  5 +++
 5 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 packages/pyright-internal/src/tests/samples/final7.py
 create mode 100644 packages/pyright-internal/src/tests/samples/final8.py

diff --git a/packages/pyright-internal/src/analyzer/checker.ts b/packages/pyright-internal/src/analyzer/checker.ts
index 0f33c72d49f9..98112e041986 100644
--- a/packages/pyright-internal/src/analyzer/checker.ts
+++ b/packages/pyright-internal/src/analyzer/checker.ts
@@ -97,7 +97,7 @@ import * as AnalyzerNodeInfo from './analyzerNodeInfo';
 import { ConstraintTracker } from './constraintTracker';
 import { getBoundCallMethod, getBoundInitMethod, getBoundNewMethod } from './constructors';
 import { addInheritedDataClassEntries } from './dataClasses';
-import { Declaration, DeclarationType, isAliasDeclaration } from './declaration';
+import { Declaration, DeclarationType, isAliasDeclaration, isVariableDeclaration } from './declaration';
 import { getNameNodeForDeclaration } from './declarationUtils';
 import { deprecatedAliases, deprecatedSpecialForms } from './deprecatedSymbols';
 import { getEnumDeclaredValueType, isEnumClassWithMembers, transformTypeForEnumMember } from './enums';
@@ -3205,8 +3205,38 @@ export class Checker extends ParseTreeWalker {
         }
     }
 
+    private _reportOverwriteOfImportedFinal(name: string, symbol: Symbol) {
+        const decls = symbol.getDeclarations();
+
+        const finalImportDecl = decls.find((decl) => {
+            if (decl.type === DeclarationType.Alias) {
+                const resolvedDecl = this._evaluator.resolveAliasDeclaration(decl, /* resolveLocalNames */ true);
+                if (resolvedDecl && isVariableDeclaration(resolvedDecl) && resolvedDecl.isFinal) {
+                    return true;
+                }
+            }
+
+            return false;
+        });
+
+        if (!finalImportDecl) {
+            return;
+        }
+
+        decls.forEach((decl) => {
+            if (decl !== finalImportDecl) {
+                this._evaluator.addDiagnostic(
+                    DiagnosticRule.reportGeneralTypeIssues,
+                    LocMessage.finalReassigned().format({ name }),
+                    getNameNodeForDeclaration(decl) ?? decl.node
+                );
+            }
+        });
+    }
+
     private _reportMultipleFinalDeclarations(name: string, symbol: Symbol, scopeType: ScopeType) {
         if (!this._evaluator.isFinalVariable(symbol)) {
+            this._reportOverwriteOfImportedFinal(name, symbol);
             return;
         }
 
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts b/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts
index fb656344fb7b..e85c84615776 100644
--- a/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts
+++ b/packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts
@@ -33,7 +33,7 @@ import { AnalyzerFileInfo } from './analyzerFileInfo';
 import { CodeFlowReferenceExpressionNode, FlowNode } from './codeFlowTypes';
 import { ConstraintTracker } from './constraintTracker';
 import { Declaration } from './declaration';
-import * as DeclarationUtils from './declarationUtils';
+import { ResolvedAliasInfo } from './declarationUtils';
 import { SymbolWithScope } from './scope';
 import { Symbol, SynthesizedTypeInfo } from './symbol';
 import { PrintTypeFlags } from './typePrinter';
@@ -684,7 +684,7 @@ export interface TypeEvaluator {
         declaration: Declaration,
         resolveLocalNames: boolean,
         options?: ResolveAliasOptions
-    ) => DeclarationUtils.ResolvedAliasInfo | undefined;
+    ) => ResolvedAliasInfo | undefined;
     getTypeOfIterable: (
         typeResult: TypeResult,
         isAsync: boolean,
diff --git a/packages/pyright-internal/src/tests/samples/final7.py b/packages/pyright-internal/src/tests/samples/final7.py
new file mode 100644
index 000000000000..275687ae3b8b
--- /dev/null
+++ b/packages/pyright-internal/src/tests/samples/final7.py
@@ -0,0 +1,8 @@
+# This sample is used in conjunction with final8.py to test that imported
+# Final symbols cannot be overwritten.
+
+from typing import Final
+
+
+var1: Final[int] = 1
+var2: Final = 2
diff --git a/packages/pyright-internal/src/tests/samples/final8.py b/packages/pyright-internal/src/tests/samples/final8.py
new file mode 100644
index 000000000000..58952ad5f4c4
--- /dev/null
+++ b/packages/pyright-internal/src/tests/samples/final8.py
@@ -0,0 +1,20 @@
+# This sample tests that final variables imported from another module
+# cannot be overwritten.
+
+from .final7 import *
+
+# This should generate an error.
+var1 = 1
+
+# This should generate an error.
+var2 = 1
+
+
+def func1():
+    from .final7 import var1, var2
+
+    # This should generate an error.
+    var1 = 1
+
+    # This should generate an error.
+    var2 = 1
diff --git a/packages/pyright-internal/src/tests/typeEvaluator4.test.ts b/packages/pyright-internal/src/tests/typeEvaluator4.test.ts
index 61a69c30485d..238266b9f4ef 100644
--- a/packages/pyright-internal/src/tests/typeEvaluator4.test.ts
+++ b/packages/pyright-internal/src/tests/typeEvaluator4.test.ts
@@ -51,6 +51,11 @@ test('Final6', () => {
     TestUtils.validateResults(analysisResults, 2);
 });
 
+test('Final8', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['final8.py']);
+    TestUtils.validateResults(analysisResults, 4);
+});
+
 test('InferredTypes1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['inferredTypes1.py']);
     TestUtils.validateResults(analysisResults, 0);