From 623a332de507ed7ff792f7f0121a6322e4fdc838 Mon Sep 17 00:00:00 2001
From: Matt Wozniski <mwozniski@bloomberg.net>
Date: Tue, 30 Aug 2022 03:04:09 -0400
Subject: [PATCH] Allow `@final` on `TypedDict`

Allow a `TypedDict` to be decorated with `@final`. Like a regular class,
mypy will emit an error if a final `TypedDict` is subclassed.

Relates-to: #7981
---
 mypy/semanal.py                     |  4 ++--
 mypy/semanal_typeddict.py           |  5 ++++-
 test-data/unit/check-typeddict.test | 23 +++++++++++++++++++++--
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/mypy/semanal.py b/mypy/semanal.py
index 3b281194d537..ca1b26d990ad 100644
--- a/mypy/semanal.py
+++ b/mypy/semanal.py
@@ -1469,8 +1469,8 @@ def analyze_typeddict_classdef(self, defn: ClassDef) -> bool:
             for decorator in defn.decorators:
                 decorator.accept(self)
                 if isinstance(decorator, RefExpr):
-                    if decorator.fullname in FINAL_DECORATOR_NAMES:
-                        self.fail("@final cannot be used with TypedDict", decorator)
+                    if decorator.fullname in FINAL_DECORATOR_NAMES and info is not None:
+                        info.is_final = True
             if info is None:
                 self.mark_incomplete(defn.name, defn)
             else:
diff --git a/mypy/semanal_typeddict.py b/mypy/semanal_typeddict.py
index 77e83e53f686..0b5b1a37a7cf 100644
--- a/mypy/semanal_typeddict.py
+++ b/mypy/semanal_typeddict.py
@@ -4,7 +4,7 @@
 
 from typing_extensions import Final
 
-from mypy import errorcodes as codes
+from mypy import errorcodes as codes, message_registry
 from mypy.errorcodes import ErrorCode
 from mypy.exprtotype import TypeTranslationError, expr_to_unanalyzed_type
 from mypy.messages import MessageBuilder
@@ -79,6 +79,9 @@ def analyze_typeddict_classdef(self, defn: ClassDef) -> tuple[bool, TypeInfo | N
                 self.api.accept(base_expr)
                 if base_expr.fullname in TPDICT_NAMES or self.is_typeddict(base_expr):
                     possible = True
+                    if isinstance(base_expr.node, TypeInfo) and base_expr.node.is_final:
+                        err = message_registry.CANNOT_INHERIT_FROM_FINAL
+                        self.fail(err.format(base_expr.node.name).value, defn, code=err.code)
         if not possible:
             return False, None
         existing_info = None
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
index 204a4e41e3f0..7fba4da071f3 100644
--- a/test-data/unit/check-typeddict.test
+++ b/test-data/unit/check-typeddict.test
@@ -2012,16 +2012,35 @@ v = {bad2: 2}  # E: Extra key "bad" for TypedDict "Value"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
-[case testCannotUseFinalDecoratorWithTypedDict]
+[case testCannotSubclassFinalTypedDict]
 from typing import TypedDict
 from typing_extensions import final
 
-@final  # E: @final cannot be used with TypedDict
+@final
 class DummyTypedDict(TypedDict):
     int_val: int
     float_val: float
     str_val: str
 
+class SubType(DummyTypedDict): # E: Cannot inherit from final class "DummyTypedDict"
+    pass
+
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.pyi]
+
+[case testCannotSubclassFinalTypedDictWithForwardDeclarations]
+from typing import TypedDict
+from typing_extensions import final
+
+@final
+class DummyTypedDict(TypedDict):
+    forward_declared: "ForwardDeclared"
+
+class SubType(DummyTypedDict): # E: Cannot inherit from final class "DummyTypedDict"
+    pass
+
+class ForwardDeclared: pass
+
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]