Skip to content

Commit e4c43cb

Browse files
sobolevnpre-commit-ci[bot]AlexWaygood
authored
[stubtest] support @type_check_only decorator (#16422)
There are several `TODO` items for the future (not in this PR): - [ ] Add an error code to disallow importing things that are decorated with `@type_check_only` - [ ] Support `@overload`ed functions. But, how? There are two options: we can treat individual overload cases as `@type_check_only` or we can treat the whole func. Since `typeshed` does not have any examples of this, I prefer to defer this discussion to somewhere else and support this when we decide Refs #15146 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Alex Waygood <[email protected]>
1 parent f7a0530 commit e4c43cb

File tree

6 files changed

+91
-0
lines changed

6 files changed

+91
-0
lines changed

mypy/nodes.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -513,6 +513,7 @@ class FuncBase(Node):
513513
"is_static", # Uses "@staticmethod" (explicit or implicit)
514514
"is_final", # Uses "@final"
515515
"is_explicit_override", # Uses "@override"
516+
"is_type_check_only", # Uses "@type_check_only"
516517
"_fullname",
517518
)
518519

@@ -530,6 +531,7 @@ def __init__(self) -> None:
530531
self.is_static = False
531532
self.is_final = False
532533
self.is_explicit_override = False
534+
self.is_type_check_only = False
533535
# Name with module prefix
534536
self._fullname = ""
535537

@@ -2866,6 +2868,7 @@ class is generic then it will be a type constructor of higher kind.
28662868
"type_var_tuple_suffix",
28672869
"self_type",
28682870
"dataclass_transform_spec",
2871+
"is_type_check_only",
28692872
)
28702873

28712874
_fullname: str # Fully qualified name
@@ -3016,6 +3019,9 @@ class is generic then it will be a type constructor of higher kind.
30163019
# Added if the corresponding class is directly decorated with `typing.dataclass_transform`
30173020
dataclass_transform_spec: DataclassTransformSpec | None
30183021

3022+
# Is set to `True` when class is decorated with `@typing.type_check_only`
3023+
is_type_check_only: bool
3024+
30193025
FLAGS: Final = [
30203026
"is_abstract",
30213027
"is_enum",
@@ -3072,6 +3078,7 @@ def __init__(self, names: SymbolTable, defn: ClassDef, module_name: str) -> None
30723078
self.metadata = {}
30733079
self.self_type = None
30743080
self.dataclass_transform_spec = None
3081+
self.is_type_check_only = False
30753082

30763083
def add_type_vars(self) -> None:
30773084
self.has_type_var_tuple_type = False

mypy/semanal.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,7 @@
251251
REVEAL_TYPE_NAMES,
252252
TPDICT_NAMES,
253253
TYPE_ALIAS_NAMES,
254+
TYPE_CHECK_ONLY_NAMES,
254255
TYPED_NAMEDTUPLE_NAMES,
255256
AnyType,
256257
CallableType,
@@ -1568,6 +1569,9 @@ def visit_decorator(self, dec: Decorator) -> None:
15681569
removed.append(i)
15691570
else:
15701571
self.fail("@final cannot be used with non-method functions", d)
1572+
elif refers_to_fullname(d, TYPE_CHECK_ONLY_NAMES):
1573+
# TODO: support `@overload` funcs.
1574+
dec.func.is_type_check_only = True
15711575
elif isinstance(d, CallExpr) and refers_to_fullname(
15721576
d.callee, DATACLASS_TRANSFORM_NAMES
15731577
):
@@ -1868,6 +1872,8 @@ def analyze_class_decorator(self, defn: ClassDef, decorator: Expression) -> None
18681872
self.fail("@runtime_checkable can only be used with protocol classes", defn)
18691873
elif decorator.fullname in FINAL_DECORATOR_NAMES:
18701874
defn.info.is_final = True
1875+
elif refers_to_fullname(decorator, TYPE_CHECK_ONLY_NAMES):
1876+
defn.info.is_type_check_only = True
18711877
elif isinstance(decorator, CallExpr) and refers_to_fullname(
18721878
decorator.callee, DATACLASS_TRANSFORM_NAMES
18731879
):

mypy/stubtest.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -484,6 +484,19 @@ def _verify_metaclass(
484484
def verify_typeinfo(
485485
stub: nodes.TypeInfo, runtime: MaybeMissing[type[Any]], object_path: list[str]
486486
) -> Iterator[Error]:
487+
if stub.is_type_check_only:
488+
# This type only exists in stubs, we only check that the runtime part
489+
# is missing. Other checks are not required.
490+
if not isinstance(runtime, Missing):
491+
yield Error(
492+
object_path,
493+
'is marked as "@type_check_only", but also exists at runtime',
494+
stub,
495+
runtime,
496+
stub_desc=repr(stub),
497+
)
498+
return
499+
487500
if isinstance(runtime, Missing):
488501
yield Error(object_path, "is not present at runtime", stub, runtime, stub_desc=repr(stub))
489502
return
@@ -1066,6 +1079,7 @@ def verify_var(
10661079
def verify_overloadedfuncdef(
10671080
stub: nodes.OverloadedFuncDef, runtime: MaybeMissing[Any], object_path: list[str]
10681081
) -> Iterator[Error]:
1082+
# TODO: support `@type_check_only` decorator
10691083
if isinstance(runtime, Missing):
10701084
yield Error(object_path, "is not present at runtime", stub, runtime)
10711085
return
@@ -1260,6 +1274,19 @@ def apply_decorator_to_funcitem(
12601274
def verify_decorator(
12611275
stub: nodes.Decorator, runtime: MaybeMissing[Any], object_path: list[str]
12621276
) -> Iterator[Error]:
1277+
if stub.func.is_type_check_only:
1278+
# This function only exists in stubs, we only check that the runtime part
1279+
# is missing. Other checks are not required.
1280+
if not isinstance(runtime, Missing):
1281+
yield Error(
1282+
object_path,
1283+
'is marked as "@type_check_only", but also exists at runtime',
1284+
stub,
1285+
runtime,
1286+
stub_desc=repr(stub),
1287+
)
1288+
return
1289+
12631290
if isinstance(runtime, Missing):
12641291
yield Error(object_path, "is not present at runtime", stub, runtime)
12651292
return

mypy/test/teststubtest.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ class Sequence(Iterable[_T_co]): ...
7171
class Tuple(Sequence[_T_co]): ...
7272
class NamedTuple(tuple[Any, ...]): ...
7373
def overload(func: _T) -> _T: ...
74+
def type_check_only(func: _T) -> _T: ...
7475
def deprecated(__msg: str) -> Callable[[_T], _T]: ...
7576
def final(func: _T) -> _T: ...
7677
"""
@@ -2046,6 +2047,50 @@ def some(self) -> int: ...
20462047
error=None,
20472048
)
20482049

2050+
@collect_cases
2051+
def test_type_check_only(self) -> Iterator[Case]:
2052+
yield Case(
2053+
stub="from typing import type_check_only, overload",
2054+
runtime="from typing import overload",
2055+
error=None,
2056+
)
2057+
# You can have public types that are only defined in stubs
2058+
# with `@type_check_only`:
2059+
yield Case(
2060+
stub="""
2061+
@type_check_only
2062+
class A1: ...
2063+
""",
2064+
runtime="",
2065+
error=None,
2066+
)
2067+
# Having `@type_check_only` on a type that exists at runtime is an error
2068+
yield Case(
2069+
stub="""
2070+
@type_check_only
2071+
class A2: ...
2072+
""",
2073+
runtime="class A2: ...",
2074+
error="A2",
2075+
)
2076+
# The same is true for functions:
2077+
yield Case(
2078+
stub="""
2079+
@type_check_only
2080+
def func1() -> None: ...
2081+
""",
2082+
runtime="",
2083+
error=None,
2084+
)
2085+
yield Case(
2086+
stub="""
2087+
@type_check_only
2088+
def func2() -> None: ...
2089+
""",
2090+
runtime="def func2() -> None: ...",
2091+
error="func2",
2092+
)
2093+
20492094

20502095
def remove_color_code(s: str) -> str:
20512096
return re.sub("\\x1b.*?m", "", s) # this works!

mypy/types.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,9 @@
113113
# Supported @final decorator names.
114114
FINAL_DECORATOR_NAMES: Final = ("typing.final", "typing_extensions.final")
115115

116+
# Supported @type_check_only names.
117+
TYPE_CHECK_ONLY_NAMES: Final = ("typing.type_check_only", "typing_extensions.type_check_only")
118+
116119
# Supported Literal type names.
117120
LITERAL_TYPE_NAMES: Final = ("typing.Literal", "typing_extensions.Literal")
118121

test-data/unit/fixtures/typing-full.pyi

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,3 +196,6 @@ def override(__arg: T) -> T: ...
196196

197197
# Was added in 3.11
198198
def reveal_type(__obj: T) -> T: ...
199+
200+
# Only exists in type checking time:
201+
def type_check_only(__func_or_class: T) -> T: ...

0 commit comments

Comments
 (0)