[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]>

Author: 3 people

mypy/nodes.py

@@ -513,6 +513,7 @@ class FuncBase(Node):
         "is_static",  # Uses "@staticmethod" (explicit or implicit)
         "is_final",  # Uses "@final"
         "is_explicit_override",  # Uses "@override"
+        "is_type_check_only",  # Uses "@type_check_only"
         "_fullname",
     )
 
@@ -530,6 +531,7 @@ def __init__(self) -> None:
         self.is_static = False
         self.is_final = False
         self.is_explicit_override = False
+        self.is_type_check_only = False
         # Name with module prefix
         self._fullname = ""
 
@@ -2866,6 +2868,7 @@ class is generic then it will be a type constructor of higher kind.
         "type_var_tuple_suffix",
         "self_type",
         "dataclass_transform_spec",
+        "is_type_check_only",
     )
 
     _fullname: str  # Fully qualified name
@@ -3016,6 +3019,9 @@ class is generic then it will be a type constructor of higher kind.
     # Added if the corresponding class is directly decorated with `typing.dataclass_transform`
     dataclass_transform_spec: DataclassTransformSpec | None
 
+    # Is set to `True` when class is decorated with `@typing.type_check_only`
+    is_type_check_only: bool
+
     FLAGS: Final = [
         "is_abstract",
         "is_enum",
@@ -3072,6 +3078,7 @@ def __init__(self, names: SymbolTable, defn: ClassDef, module_name: str) -> None
         self.metadata = {}
         self.self_type = None
         self.dataclass_transform_spec = None
+        self.is_type_check_only = False
 
     def add_type_vars(self) -> None:
         self.has_type_var_tuple_type = False

mypy/semanal.py

@@ -251,6 +251,7 @@
     REVEAL_TYPE_NAMES,
     TPDICT_NAMES,
     TYPE_ALIAS_NAMES,
+    TYPE_CHECK_ONLY_NAMES,
     TYPED_NAMEDTUPLE_NAMES,
     AnyType,
     CallableType,
@@ -1568,6 +1569,9 @@ def visit_decorator(self, dec: Decorator) -> None:
                     removed.append(i)
                 else:
                     self.fail("@final cannot be used with non-method functions", d)
+            elif refers_to_fullname(d, TYPE_CHECK_ONLY_NAMES):
+                # TODO: support `@overload` funcs.
+                dec.func.is_type_check_only = True
             elif isinstance(d, CallExpr) and refers_to_fullname(
                 d.callee, DATACLASS_TRANSFORM_NAMES
             ):
@@ -1868,6 +1872,8 @@ def analyze_class_decorator(self, defn: ClassDef, decorator: Expression) -> None
                     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
         ):

mypy/stubtest.py

@@ -484,6 +484,19 @@ def _verify_metaclass(
 def verify_typeinfo(
     stub: nodes.TypeInfo, runtime: MaybeMissing[type[Any]], object_path: list[str]
 ) -> Iterator[Error]:
+    if stub.is_type_check_only:
+        # This type only exists in stubs, we only check that the runtime part
+        # is missing. Other checks are not required.
+        if not isinstance(runtime, Missing):
+            yield Error(
+                object_path,
+                'is marked as "@type_check_only", but also exists at runtime',
+                stub,
+                runtime,
+                stub_desc=repr(stub),
+            )
+        return
+
     if isinstance(runtime, Missing):
         yield Error(object_path, "is not present at runtime", stub, runtime, stub_desc=repr(stub))
         return
@@ -1066,6 +1079,7 @@ def verify_var(
 def verify_overloadedfuncdef(
     stub: nodes.OverloadedFuncDef, runtime: MaybeMissing[Any], object_path: list[str]
 ) -> Iterator[Error]:
+    # TODO: support `@type_check_only` decorator
     if isinstance(runtime, Missing):
         yield Error(object_path, "is not present at runtime", stub, runtime)
         return
@@ -1260,6 +1274,19 @@ def apply_decorator_to_funcitem(
 def verify_decorator(
     stub: nodes.Decorator, runtime: MaybeMissing[Any], object_path: list[str]
 ) -> Iterator[Error]:
+    if stub.func.is_type_check_only:
+        # This function only exists in stubs, we only check that the runtime part
+        # is missing. Other checks are not required.
+        if not isinstance(runtime, Missing):
+            yield Error(
+                object_path,
+                'is marked as "@type_check_only", but also exists at runtime',
+                stub,
+                runtime,
+                stub_desc=repr(stub),
+            )
+        return
+
     if isinstance(runtime, Missing):
         yield Error(object_path, "is not present at runtime", stub, runtime)
         return

mypy/test/teststubtest.py

@@ -71,6 +71,7 @@ class Sequence(Iterable[_T_co]): ...
 class Tuple(Sequence[_T_co]): ...
 class NamedTuple(tuple[Any, ...]): ...
 def overload(func: _T) -> _T: ...
+def type_check_only(func: _T) -> _T: ...
 def deprecated(__msg: str) -> Callable[[_T], _T]: ...
 def final(func: _T) -> _T: ...
 """
@@ -2046,6 +2047,50 @@ def some(self) -> int: ...
             error=None,
         )
 
+    @collect_cases
+    def test_type_check_only(self) -> Iterator[Case]:
+        yield Case(
+            stub="from typing import type_check_only, overload",
+            runtime="from typing import overload",
+            error=None,
+        )
+        # You can have public types that are only defined in stubs
+        # with `@type_check_only`:
+        yield Case(
+            stub="""
+            @type_check_only
+            class A1: ...
+            """,
+            runtime="",
+            error=None,
+        )
+        # Having `@type_check_only` on a type that exists at runtime is an error
+        yield Case(
+            stub="""
+            @type_check_only
+            class A2: ...
+            """,
+            runtime="class A2: ...",
+            error="A2",
+        )
+        # The same is true for functions:
+        yield Case(
+            stub="""
+            @type_check_only
+            def func1() -> None: ...
+            """,
+            runtime="",
+            error=None,
+        )
+        yield Case(
+            stub="""
+            @type_check_only
+            def func2() -> None: ...
+            """,
+            runtime="def func2() -> None: ...",
+            error="func2",
+        )
+
 
 def remove_color_code(s: str) -> str:
     return re.sub("\\x1b.*?m", "", s)  # this works!

mypy/types.py

@@ -113,6 +113,9 @@
 # Supported @final decorator names.
 FINAL_DECORATOR_NAMES: Final = ("typing.final", "typing_extensions.final")
 
+# Supported @type_check_only names.
+TYPE_CHECK_ONLY_NAMES: Final = ("typing.type_check_only", "typing_extensions.type_check_only")
+
 # Supported Literal type names.
 LITERAL_TYPE_NAMES: Final = ("typing.Literal", "typing_extensions.Literal")
 

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

@@ -196,3 +196,6 @@ def override(__arg: T) -> T: ...
 
 # Was added in 3.11
 def reveal_type(__obj: T) -> T: ...
+
+# Only exists in type checking time:
+def type_check_only(__func_or_class: T) -> T: ...