Refactor class decorator: this enables type_check_only support for TypedDict and NamedTuple

mypy/semanal.py

@@ -1743,9 +1743,8 @@ def analyze_typeddict_classdef(self, defn: ClassDef) -> bool:
         if is_typeddict:
             for decorator in defn.decorators:
                 decorator.accept(self)
-                if isinstance(decorator, RefExpr):
-                    if decorator.fullname in FINAL_DECORATOR_NAMES and info is not None:
-                        info.is_final = True
+                if info is not None:
+                    self.analyze_class_decorator_common(defn, info, decorator)
             if info is None:
                 self.mark_incomplete(defn.name, defn)
             else:
@@ -1781,8 +1780,7 @@ def analyze_namedtuple_classdef(
                 with self.scope.class_scope(defn.info):
                     for deco in defn.decorators:
                         deco.accept(self)
-                        if isinstance(deco, RefExpr) and deco.fullname in FINAL_DECORATOR_NAMES:
-                            info.is_final = True
+                        self.analyze_class_decorator_common(defn, defn.info, deco)
                     with self.named_tuple_analyzer.save_namedtuple_body(info):
                         self.analyze_class_body_common(defn)
             return True
@@ -1864,21 +1862,30 @@ def leave_class(self) -> None:
 
     def analyze_class_decorator(self, defn: ClassDef, decorator: Expression) -> None:
         decorator.accept(self)
+        self.analyze_class_decorator_common(defn, defn.info, decorator)
         if isinstance(decorator, RefExpr):
             if decorator.fullname in RUNTIME_PROTOCOL_DECOS:
                 if defn.info.is_protocol:
                     defn.info.runtime_protocol = True
                 else:
                     self.fail("@runtime_checkable can only be used with protocol classes", defn)
-            elif decorator.fullname in FINAL_DECORATOR_NAMES:
-                defn.info.is_final = True
-            elif refers_to_fullname(decorator, TYPE_CHECK_ONLY_NAMES):
-                defn.info.is_type_check_only = True
         elif isinstance(decorator, CallExpr) and refers_to_fullname(
             decorator.callee, DATACLASS_TRANSFORM_NAMES
         ):
             defn.info.dataclass_transform_spec = self.parse_dataclass_transform_spec(decorator)
 
+    def analyze_class_decorator_common(
+        self, defn: ClassDef, info: TypeInfo, decorator: Expression
+    ) -> None:
+        """Common part for all class decorators.
+
+        Including classes, typeddicts, and namedtuples.
+        """
+        if refers_to_fullname(decorator, FINAL_DECORATOR_NAMES):
+            info.is_final = True
+        elif refers_to_fullname(decorator, TYPE_CHECK_ONLY_NAMES):
+            info.is_type_check_only = True
+
     def clean_up_bases_and_infer_type_variables(
         self, defn: ClassDef, base_type_exprs: list[Expression], context: Context
     ) -> tuple[list[Expression], list[TypeVarLikeType], bool]:

mypy/test/teststubtest.py

@@ -2073,6 +2073,28 @@ class A2: ...
             runtime="class A2: ...",
             error="A2",
         )
+        # The same is true for NamedTuples and TypedDicts:
+        yield Case(
+            stub="from typing_extensions import NamedTuple, TypedDict",
+            runtime="from typing_extensions import NamedTuple, TypedDict",
+            error=None,
+        )
+        yield Case(
+            stub="""
+            @type_check_only
+            class NT1(NamedTuple): ...
+            """,
+            runtime="class NT1(NamedTuple): ...",
+            error="NT1",
+        )
+        yield Case(
+            stub="""
+            @type_check_only
+            class TD1(TypedDict): ...
+            """,
+            runtime="class TD1(TypedDict): ...",
+            error="TD1",
+        )
         # The same is true for functions:
         yield Case(
             stub="""