From e4c43cb68b742f0f51759565d1d3c4a722d16f55 Mon Sep 17 00:00:00 2001 From: Nikita Sobolev Date: Sun, 12 Nov 2023 00:00:11 +0300 Subject: [PATCH] [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 https://github.com/python/mypy/issues/15146 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Alex Waygood --- mypy/nodes.py | 7 ++++ mypy/semanal.py | 6 ++++ mypy/stubtest.py | 27 +++++++++++++++ mypy/test/teststubtest.py | 45 +++++++++++++++++++++++++ mypy/types.py | 3 ++ test-data/unit/fixtures/typing-full.pyi | 3 ++ 6 files changed, 91 insertions(+) diff --git a/mypy/nodes.py b/mypy/nodes.py index d65a23a6b7fe..17e06613d1e3 100644 --- a/mypy/nodes.py +++ b/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 diff --git a/mypy/semanal.py b/mypy/semanal.py index 6f322af816ea..68f0d04e77ca 100644 --- a/mypy/semanal.py +++ b/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 ): diff --git a/mypy/stubtest.py b/mypy/stubtest.py index adffe2003ad4..46468e8e18e0 100644 --- a/mypy/stubtest.py +++ b/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 diff --git a/mypy/test/teststubtest.py b/mypy/test/teststubtest.py index 0c1817202f1f..6a973d16d7bc 100644 --- a/mypy/test/teststubtest.py +++ b/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! diff --git a/mypy/types.py b/mypy/types.py index b100cf569086..e7738bd7d088 100644 --- a/mypy/types.py +++ b/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") diff --git a/test-data/unit/fixtures/typing-full.pyi b/test-data/unit/fixtures/typing-full.pyi index ef903ace78af..ca8a2413f05f 100644 --- a/test-data/unit/fixtures/typing-full.pyi +++ b/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: ...