Skip to content

Commit 023eb41

Browse files
authored
Refactor class decorator: this enables type_check_only support for TypedDict and NamedTuple (#16469)
I've noticed that `TypedDict` and `NamedTuple` classes are special cased during semantic analyzisys. They had their own logic for class-level decorators. This is fine, but we need some common ground. As a side-effect, they can now be `type_check_only`!
1 parent efa5dcb commit 023eb41

File tree

2 files changed

+38
-9
lines changed

2 files changed

+38
-9
lines changed

mypy/semanal.py

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1743,9 +1743,8 @@ def analyze_typeddict_classdef(self, defn: ClassDef) -> bool:
17431743
if is_typeddict:
17441744
for decorator in defn.decorators:
17451745
decorator.accept(self)
1746-
if isinstance(decorator, RefExpr):
1747-
if decorator.fullname in FINAL_DECORATOR_NAMES and info is not None:
1748-
info.is_final = True
1746+
if info is not None:
1747+
self.analyze_class_decorator_common(defn, info, decorator)
17491748
if info is None:
17501749
self.mark_incomplete(defn.name, defn)
17511750
else:
@@ -1781,8 +1780,7 @@ def analyze_namedtuple_classdef(
17811780
with self.scope.class_scope(defn.info):
17821781
for deco in defn.decorators:
17831782
deco.accept(self)
1784-
if isinstance(deco, RefExpr) and deco.fullname in FINAL_DECORATOR_NAMES:
1785-
info.is_final = True
1783+
self.analyze_class_decorator_common(defn, defn.info, deco)
17861784
with self.named_tuple_analyzer.save_namedtuple_body(info):
17871785
self.analyze_class_body_common(defn)
17881786
return True
@@ -1864,21 +1862,30 @@ def leave_class(self) -> None:
18641862

18651863
def analyze_class_decorator(self, defn: ClassDef, decorator: Expression) -> None:
18661864
decorator.accept(self)
1865+
self.analyze_class_decorator_common(defn, defn.info, decorator)
18671866
if isinstance(decorator, RefExpr):
18681867
if decorator.fullname in RUNTIME_PROTOCOL_DECOS:
18691868
if defn.info.is_protocol:
18701869
defn.info.runtime_protocol = True
18711870
else:
18721871
self.fail("@runtime_checkable can only be used with protocol classes", defn)
1873-
elif decorator.fullname in FINAL_DECORATOR_NAMES:
1874-
defn.info.is_final = True
1875-
elif refers_to_fullname(decorator, TYPE_CHECK_ONLY_NAMES):
1876-
defn.info.is_type_check_only = True
18771872
elif isinstance(decorator, CallExpr) and refers_to_fullname(
18781873
decorator.callee, DATACLASS_TRANSFORM_NAMES
18791874
):
18801875
defn.info.dataclass_transform_spec = self.parse_dataclass_transform_spec(decorator)
18811876

1877+
def analyze_class_decorator_common(
1878+
self, defn: ClassDef, info: TypeInfo, decorator: Expression
1879+
) -> None:
1880+
"""Common method for applying class decorators.
1881+
1882+
Called on regular classes, typeddicts, and namedtuples.
1883+
"""
1884+
if refers_to_fullname(decorator, FINAL_DECORATOR_NAMES):
1885+
info.is_final = True
1886+
elif refers_to_fullname(decorator, TYPE_CHECK_ONLY_NAMES):
1887+
info.is_type_check_only = True
1888+
18821889
def clean_up_bases_and_infer_type_variables(
18831890
self, defn: ClassDef, base_type_exprs: list[Expression], context: Context
18841891
) -> tuple[list[Expression], list[TypeVarLikeType], bool]:

mypy/test/teststubtest.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2073,6 +2073,28 @@ class A2: ...
20732073
runtime="class A2: ...",
20742074
error="A2",
20752075
)
2076+
# The same is true for NamedTuples and TypedDicts:
2077+
yield Case(
2078+
stub="from typing_extensions import NamedTuple, TypedDict",
2079+
runtime="from typing_extensions import NamedTuple, TypedDict",
2080+
error=None,
2081+
)
2082+
yield Case(
2083+
stub="""
2084+
@type_check_only
2085+
class NT1(NamedTuple): ...
2086+
""",
2087+
runtime="class NT1(NamedTuple): ...",
2088+
error="NT1",
2089+
)
2090+
yield Case(
2091+
stub="""
2092+
@type_check_only
2093+
class TD1(TypedDict): ...
2094+
""",
2095+
runtime="class TD1(TypedDict): ...",
2096+
error="TD1",
2097+
)
20762098
# The same is true for functions:
20772099
yield Case(
20782100
stub="""

0 commit comments

Comments
 (0)