Skip to content

Commit

Permalink
Refactor class decorator: this enables type_check_only support for …
Browse files Browse the repository at this point in the history
…`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`!
  • Loading branch information
sobolevn authored Nov 13, 2023
1 parent efa5dcb commit 023eb41
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 9 deletions.
25 changes: 16 additions & 9 deletions mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 method for applying class decorators.
Called on regular 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]:
Expand Down
22 changes: 22 additions & 0 deletions mypy/test/teststubtest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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="""
Expand Down

0 comments on commit 023eb41

Please sign in to comment.