From 613db60dc27684a47646642b9292648e195faadf Mon Sep 17 00:00:00 2001 From: Stephen Moore Date: Tue, 28 May 2024 14:32:29 +1000 Subject: [PATCH 1/5] Extracting the Known Annotations enum into a separate file --- .../plugin/_known_annotations.py | 7 +++++++ extended_mypy_django_plugin/plugin/_plugin.py | 17 +++++++++++------ 2 files changed, 18 insertions(+), 6 deletions(-) create mode 100644 extended_mypy_django_plugin/plugin/_known_annotations.py diff --git a/extended_mypy_django_plugin/plugin/_known_annotations.py b/extended_mypy_django_plugin/plugin/_known_annotations.py new file mode 100644 index 0000000..2a00f76 --- /dev/null +++ b/extended_mypy_django_plugin/plugin/_known_annotations.py @@ -0,0 +1,7 @@ +import enum + + +class KnownAnnotations(enum.Enum): + CONCRETE = "extended_mypy_django_plugin.annotations.Concrete" + CONCRETE_QUERYSET = "extended_mypy_django_plugin.annotations.ConcreteQuerySet" + DEFAULT_QUERYSET = "extended_mypy_django_plugin.annotations.DefaultQuerySet" diff --git a/extended_mypy_django_plugin/plugin/_plugin.py b/extended_mypy_django_plugin/plugin/_plugin.py index 1611314..b4611e2 100644 --- a/extended_mypy_django_plugin/plugin/_plugin.py +++ b/extended_mypy_django_plugin/plugin/_plugin.py @@ -1,5 +1,4 @@ import collections -import enum import sys from typing import Generic @@ -26,7 +25,16 @@ ) from typing_extensions import assert_never -from . import _config, _dependencies, _helpers, _hook, _reports, _store, actions +from . import ( + _config, + _dependencies, + _helpers, + _hook, + _known_annotations, + _reports, + _store, + actions, +) class Hook( @@ -62,10 +70,7 @@ class ExtendedMypyStubs(main.NewSemanalDjangoPlugin): plugin_config: _config.Config - class Annotations(enum.Enum): - CONCRETE = "extended_mypy_django_plugin.annotations.Concrete" - CONCRETE_QUERYSET = "extended_mypy_django_plugin.annotations.ConcreteQuerySet" - DEFAULT_QUERYSET = "extended_mypy_django_plugin.annotations.DefaultQuerySet" + Annotations = _known_annotations.KnownAnnotations def __init__(self, options: Options, mypy_version_tuple: tuple[int, int]) -> None: super(main.NewSemanalDjangoPlugin, self).__init__(options) From 64f6d7c9db2a380a98e989600154ab36cee58412 Mon Sep 17 00:00:00 2001 From: Stephen Moore Date: Tue, 28 May 2024 14:38:17 +1000 Subject: [PATCH 2/5] Make sure we resolve all concrete type annotations when using type var --- extended_mypy_django_plugin/plugin/_plugin.py | 21 +- .../plugin/actions/_type_analyze.py | 2 +- .../plugin/actions/_type_checker.py | 309 ++++++++++++++---- 3 files changed, 268 insertions(+), 64 deletions(-) diff --git a/extended_mypy_django_plugin/plugin/_plugin.py b/extended_mypy_django_plugin/plugin/_plugin.py index b4611e2..a2921f4 100644 --- a/extended_mypy_django_plugin/plugin/_plugin.py +++ b/extended_mypy_django_plugin/plugin/_plugin.py @@ -276,11 +276,13 @@ def run(self, ctx: AnalyzeTypeContext) -> MypyType: @_hook.hook class get_function_hook(Hook[FunctionContext, MypyType]): """ - Find functions that return a ``DefaultQuerySet`` annotation with a type variable - and resolve the annotation. + Find functions that return a concrete annotation with a type var and resolve that annotation """ def choose(self) -> bool: + if self.fullname.startswith("builtins."): + return False + sym = self.plugin.lookup_fully_qualified(self.fullname) if not sym or not sym.node: return False @@ -294,13 +296,14 @@ def choose(self) -> bool: def run(self, ctx: FunctionContext) -> MypyType: assert isinstance(ctx.api, TypeChecker) - type_checking = actions.TypeChecking(self.store, api=ctx.api) - - result = type_checking.modify_default_queryset_return_type( - ctx, - desired_annotation_fullname=ExtendedMypyStubs.Annotations.DEFAULT_QUERYSET.value, + type_checking = actions.TypeChecking( + self.store, + api=ctx.api, + lookup_info=self.plugin._lookup_info, ) + result = type_checking.modify_return_type(ctx) + if result is not None: return result @@ -322,7 +325,9 @@ def choose(self) -> bool: def run(self, ctx: AttributeContext) -> MypyType: assert isinstance(ctx.api, TypeChecker) - type_checking = actions.TypeChecking(self.store, api=ctx.api) + type_checking = actions.TypeChecking( + self.store, api=ctx.api, lookup_info=self.plugin._lookup_info + ) return type_checking.extended_get_attribute_resolve_manager_method( ctx, resolve_manager_method_from_instance=resolve_manager_method_from_instance diff --git a/extended_mypy_django_plugin/plugin/actions/_type_analyze.py b/extended_mypy_django_plugin/plugin/actions/_type_analyze.py index 9a0912c..1017e6b 100644 --- a/extended_mypy_django_plugin/plugin/actions/_type_analyze.py +++ b/extended_mypy_django_plugin/plugin/actions/_type_analyze.py @@ -30,7 +30,7 @@ def find_concrete_models(self, unbound_type: UnboundType) -> MypyType: type_arg = get_proper_type(self.api.analyze_type(args[0])) if not isinstance(type_arg, Instance): - return UnionType(()) + return unbound_type concrete = tuple( self.store.retrieve_concrete_children_types( diff --git a/extended_mypy_django_plugin/plugin/actions/_type_checker.py b/extended_mypy_django_plugin/plugin/actions/_type_checker.py index dc26782..dffa596 100644 --- a/extended_mypy_django_plugin/plugin/actions/_type_checker.py +++ b/extended_mypy_django_plugin/plugin/actions/_type_checker.py @@ -1,16 +1,15 @@ +import dataclasses +from collections.abc import Iterator, Mapping from typing import Protocol from mypy.checker import TypeChecker -from mypy.nodes import CallExpr, MemberExpr, TypeInfo -from mypy.plugin import ( - AttributeContext, - FunctionContext, -) +from mypy.nodes import CallExpr, Context, MemberExpr, TypeInfo +from mypy.plugin import AttributeContext, FunctionContext, MethodContext from mypy.types import ( AnyType, CallableType, - FormalArgument, Instance, + ProperType, TypeOfAny, TypeType, TypeVarType, @@ -19,8 +18,13 @@ get_proper_type, ) from mypy.types import Type as MypyType +from typing_extensions import Self + +from .. import _known_annotations, _store -from .. import _store + +class LookupFunction(Protocol): + def __call__(self, fullname: str) -> TypeInfo | None: ... class ResolveManagerMethodFromInstance(Protocol): @@ -29,21 +33,146 @@ def __call__( ) -> MypyType: ... +def _find_type_vars( + item: MypyType, _chain: list[ProperType] | None = None +) -> list[TypeVarType | str]: + if _chain is None: + _chain = [] + + result: list[TypeVarType | str] = [] + + item = get_proper_type(item) + + if isinstance(item, TypeVarType): + result.append(item) + + elif isinstance(item, UnboundType): + if item.args: + for arg in item.args: + if arg not in _chain: + _chain.append(get_proper_type(arg)) + result.extend(_find_type_vars(arg, _chain=_chain)) + else: + _chain.append(item) + result.append(item.name) + + return result + + +@dataclasses.dataclass +class BasicTypeInfo: + is_type: bool + api: TypeChecker + func: CallableType + + item: ProperType + type_vars: list[TypeVarType | str] + concrete_annotation: _known_annotations.KnownAnnotations | None + + @classmethod + def create(cls, api: TypeChecker, func: CallableType, item: MypyType | None = None) -> Self: + is_type: bool = False + + if item is None: + item = func.ret_type + + item = get_proper_type(item) + if isinstance(item, TypeType): + is_type = True + item = item.item + + type_vars = _find_type_vars(item) + + concrete_annotation: _known_annotations.KnownAnnotations | None = None + if isinstance(item, UnboundType): + try: + named_generic_type_name = api.named_generic_type( + item.name, list(item.args) + ).type.fullname + except AssertionError: + named_generic_type_name = "" + + try: + concrete_annotation = _known_annotations.KnownAnnotations(named_generic_type_name) + except ValueError: + pass + + return cls( + api=api, + func=func, + item=item, + is_type=is_type, + type_vars=type_vars, + concrete_annotation=concrete_annotation, + ) + + @property + def contains_concrete_annotation(self) -> bool: + if self.concrete_annotation is not None: + return True + + for item in self.items(): + if item is not self and item.contains_concrete_annotation: + return True + + return False + + def items(self) -> Iterator[Self]: + if isinstance(self.item, UnionType): + for item in self.item.items: + yield self.__class__.create(self.api, self.func, item) + else: + yield self + + def map_type_vars( + self, context: Context, callee_arg_names: list[str | None], arg_types: list[list[MypyType]] + ) -> Mapping[TypeVarType | str, ProperType]: + result: dict[TypeVarType | str, ProperType] = {} + + formal_by_name = {arg.name: arg.typ for arg in self.func.formal_arguments()} + + for arg_name, arg_type in zip(callee_arg_names, arg_types): + underlying = get_proper_type(formal_by_name[arg_name]) + if isinstance(underlying, TypeType): + underlying = underlying.item + + if isinstance(underlying, TypeVarType): + found_type = get_proper_type(arg_type[0]) + if isinstance(found_type, CallableType): + found_type = get_proper_type(found_type.ret_type) + + result[underlying] = found_type + result[underlying.name] = found_type + + for type_var in self.type_vars: + if type_var not in result: + self.api.fail( + f"Failed to find an argument that matched the type var {type_var}", context + ) + result[type_var] = AnyType(TypeOfAny.from_error) + + return result + + class TypeChecking: - def __init__(self, store: _store.Store, *, api: TypeChecker) -> None: + def __init__( + self, store: _store.Store, *, api: TypeChecker, lookup_info: LookupFunction + ) -> None: self.api = api self.store = store + self._lookup_info = lookup_info - def modify_default_queryset_return_type( - self, - ctx: FunctionContext, - *, - desired_annotation_fullname: str, - ) -> MypyType | None: - if not isinstance(ctx.default_return_type, UnboundType): - return None + def _named_type(self, fullname: str, args: list[MypyType] | None = None) -> Instance: + """ + Copied from what semantic analyzer does + """ + node = self._lookup_info(fullname) + assert isinstance(node, TypeInfo) + if args: + return Instance(node, args) + return Instance(node, [AnyType(TypeOfAny.special_form)] * len(node.defn.type_vars)) - context = ctx.context + def _get_info(self, context: Context) -> BasicTypeInfo | None: if not isinstance(context, CallExpr): return None @@ -58,59 +187,119 @@ def modify_default_queryset_return_type( if not isinstance(func, CallableType): return None - if not isinstance(func.ret_type, UnboundType): + return BasicTypeInfo.create(api=self.api, func=func) + + def modify_return_type(self, ctx: MethodContext | FunctionContext) -> MypyType | None: + info = self._get_info(ctx.context) + if info is None: return None - as_generic_type = self.api.named_generic_type(func.ret_type.name, [func.ret_type.args[0]]) - if as_generic_type.type.fullname != desired_annotation_fullname: + if not info.contains_concrete_annotation: return None - if len(func.ret_type.args) != 1: - self.api.fail("DefaultQuerySet takes only one argument", context) - return AnyType(TypeOfAny.from_error) + result: list[MypyType] = [] - found_type: MypyType | None = None + type_vars_map = info.map_type_vars(ctx.context, ctx.callee_arg_names, ctx.arg_types) - type_var = func.ret_type.args[0] + for item in info.items(): + Known = _known_annotations.KnownAnnotations + if item.concrete_annotation is None: + if isinstance(item.item, TypeVarType): + result.append(type_vars_map.get(item.item, item.item)) + elif isinstance(item.item, UnboundType) and len(item.item.args) == 0: + result.append(type_vars_map.get(item.item.name, item.item)) + else: + result.append(item.item) + continue - if isinstance(type_var, UnboundType): - match: FormalArgument | None = None - for arg in func.formal_arguments(): - arg_name: str | None = None - formal_arg_type = get_proper_type(arg.typ) - if isinstance(formal_arg_type, TypeType) and isinstance( - formal_arg_type.item, TypeVarType - ): - arg_name = formal_arg_type.item.name + is_type = item.is_type and not info.is_type - elif isinstance(arg, TypeVarType): - arg_name = formal_arg_type.name + # It has to be an instance or unbound type if it has a concrete annotation + assert isinstance(item.item, Instance | UnboundType) - if arg_name and arg_name == type_var.name: - match = arg - break + if len(item.item.args) != 1: + self.api.fail("Concrete Annotations must take exactly one argument", ctx.context) + return AnyType(TypeOfAny.from_error) - if match is not None: - for arg_name, arg_type in zip(ctx.callee_arg_names, ctx.arg_types): - if arg_name == match.name: - found_type = arg_type[0] + model = item.item.args[0] + if isinstance(model, TypeVarType | UnboundType): + found: ProperType | None = None + if isinstance(model, TypeVarType): + found = type_vars_map.get(model) + elif isinstance(model, UnboundType): + found = type_vars_map.get(model.name) - if found_type is None: - self.api.fail("Failed to find an argument that matched the type var", context) - return AnyType(TypeOfAny.from_error) + if found is None: + self.api.fail( + f"Can't determine what model the type var {model} represents", ctx.context + ) + return AnyType(TypeOfAny.from_error) + else: + if isinstance(found, AnyType): + return found - found_type = get_proper_type(found_type) - if isinstance(found_type, CallableType): - type_var = found_type.ret_type - else: - type_var = found_type + model = found - if not isinstance(type_var, Instance | UnionType): - self.api.fail("Don't know what to do with what DefaultQuerySet was given", context) - return AnyType(TypeOfAny.from_error) + model = get_proper_type(model) + instances: list[Instance] = [] + if isinstance(model, Instance): + instances.append(model) + elif isinstance(model, UnionType): + for member in model.items: + member = get_proper_type(member) + if isinstance(member, Instance): + instances.append(member) + else: + self.api.fail( + f"Failed to have a list of instances to find for a concrete annotation, {member}", + ctx.context, + ) + return AnyType(TypeOfAny.from_error) + made: MypyType + if item.concrete_annotation is Known.CONCRETE: + made = self.get_concrete_types(ctx.context, instances=instances) + elif item.concrete_annotation is Known.DEFAULT_QUERYSET: + made = self.get_default_queryset_return_type( + ctx.context, instances=UnionType(tuple(instances)) + ) + elif item.concrete_annotation is Known.CONCRETE_QUERYSET: + made = self.get_concrete_queryset_return_type(ctx.context, instances=instances) + + if is_type: + made = TypeType(made) + + result.append(made) + + final: MypyType = UnionType(tuple(result)) + if info.is_type: + final = TypeType(final) + + return final + + def get_concrete_queryset_return_type( + self, context: Context, *, instances: list[Instance] + ) -> MypyType: + result: list[MypyType] = [] + for concrete in self.get_concrete_types(context, instances=instances).items: + concrete = get_proper_type(concrete) + assert isinstance(concrete, Instance) + try: + result.extend(self.store.realise_querysets(concrete, self.lookup_info)) + except _store.RestartDmypy: + self.api.fail("You probably need to restart dmypy", context) + return AnyType(TypeOfAny.from_error) + except _store.UnionMustBeOfTypes: + self.api.fail("Union must be of instances of models", context) + return AnyType(TypeOfAny.from_error) + + return UnionType(tuple(result)) + + def get_default_queryset_return_type( + self, context: Context, *, instances: Instance | UnionType + ) -> MypyType: try: - querysets = tuple(self.store.realise_querysets(type_var, self.lookup_info)) + querysets = tuple(self.store.realise_querysets(instances, self.lookup_info)) except _store.RestartDmypy: self.api.fail("You probably need to restart dmypy", context) return AnyType(TypeOfAny.from_error) @@ -120,6 +309,16 @@ def modify_default_queryset_return_type( else: return UnionType(querysets) + def get_concrete_types(self, context: Context, *, instances: list[Instance]) -> UnionType: + result: list[MypyType] = [] + for item in instances: + result.extend( + self.store.retrieve_concrete_children_types( + item.type, self.lookup_info, self._named_type + ) + ) + return UnionType(tuple(result)) + def extended_get_attribute_resolve_manager_method( self, ctx: AttributeContext, From 58acb6252c65364f848aeee44d7bf17eafe6cf59 Mon Sep 17 00:00:00 2001 From: Stephen Moore Date: Tue, 28 May 2024 14:39:11 +1000 Subject: [PATCH 3/5] Move get_function_hook to the bottom of the plugin --- extended_mypy_django_plugin/plugin/_plugin.py | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/extended_mypy_django_plugin/plugin/_plugin.py b/extended_mypy_django_plugin/plugin/_plugin.py index a2921f4..1c82699 100644 --- a/extended_mypy_django_plugin/plugin/_plugin.py +++ b/extended_mypy_django_plugin/plugin/_plugin.py @@ -273,6 +273,27 @@ def run(self, ctx: AnalyzeTypeContext) -> MypyType: return method(unbound_type=ctx.type) + @_hook.hook + class get_attribute_hook(Hook[AttributeContext, MypyType]): + """ + An implementation of the change found in + https://github.com/typeddjango/django-stubs/pull/2027 + """ + + def choose(self) -> bool: + return self.super_hook is resolve_manager_method + + def run(self, ctx: AttributeContext) -> MypyType: + assert isinstance(ctx.api, TypeChecker) + + type_checking = actions.TypeChecking( + self.store, api=ctx.api, lookup_info=self.plugin._lookup_info + ) + + return type_checking.extended_get_attribute_resolve_manager_method( + ctx, resolve_manager_method_from_instance=resolve_manager_method_from_instance + ) + @_hook.hook class get_function_hook(Hook[FunctionContext, MypyType]): """ @@ -311,24 +332,3 @@ def run(self, ctx: FunctionContext) -> MypyType: return self.super_hook(ctx) return ctx.default_return_type - - @_hook.hook - class get_attribute_hook(Hook[AttributeContext, MypyType]): - """ - An implementation of the change found in - https://github.com/typeddjango/django-stubs/pull/2027 - """ - - def choose(self) -> bool: - return self.super_hook is resolve_manager_method - - def run(self, ctx: AttributeContext) -> MypyType: - assert isinstance(ctx.api, TypeChecker) - - type_checking = actions.TypeChecking( - self.store, api=ctx.api, lookup_info=self.plugin._lookup_info - ) - - return type_checking.extended_get_attribute_resolve_manager_method( - ctx, resolve_manager_method_from_instance=resolve_manager_method_from_instance - ) From db1ef7156bde4e218aa3288c15f484e360ece2bc Mon Sep 17 00:00:00 2001 From: Stephen Moore Date: Tue, 28 May 2024 14:41:10 +1000 Subject: [PATCH 4/5] Allow modifying method return types as well as function return types --- extended_mypy_django_plugin/plugin/_plugin.py | 57 +++++++++++++++++-- 1 file changed, 51 insertions(+), 6 deletions(-) diff --git a/extended_mypy_django_plugin/plugin/_plugin.py b/extended_mypy_django_plugin/plugin/_plugin.py index 1c82699..e1268a1 100644 --- a/extended_mypy_django_plugin/plugin/_plugin.py +++ b/extended_mypy_django_plugin/plugin/_plugin.py @@ -12,6 +12,7 @@ ClassDefContext, DynamicClassDefContext, FunctionContext, + MethodContext, ) from mypy.semanal import SemanticAnalyzer from mypy.typeanal import TypeAnalyser @@ -294,13 +295,24 @@ def run(self, ctx: AttributeContext) -> MypyType: ctx, resolve_manager_method_from_instance=resolve_manager_method_from_instance ) - @_hook.hook - class get_function_hook(Hook[FunctionContext, MypyType]): + class SharedCallableHookLogic: """ - Find functions that return a concrete annotation with a type var and resolve that annotation + Shared logic for modifying the return type of methods and functions that use a concrete + annotation with a type variable. """ + def __init__(self, fullname: str, plugin: "ExtendedMypyStubs") -> None: + self.plugin = plugin + self.store = plugin.store + self.fullname = fullname + def choose(self) -> bool: + """ + Choose methods and functions either returning a type guard or have a generic + return type. + + We determine whether the return type is a concrete annotation or not in the run method. + """ if self.fullname.startswith("builtins."): return False @@ -312,9 +324,9 @@ def choose(self) -> bool: if not isinstance(call, CallableType): return False - return call.is_generic() + return bool(call.type_guard or call.is_generic()) - def run(self, ctx: FunctionContext) -> MypyType: + def run(self, ctx: MethodContext | FunctionContext) -> MypyType | None: assert isinstance(ctx.api, TypeChecker) type_checking = actions.TypeChecking( @@ -323,7 +335,40 @@ def run(self, ctx: FunctionContext) -> MypyType: lookup_info=self.plugin._lookup_info, ) - result = type_checking.modify_return_type(ctx) + return type_checking.modify_return_type(ctx) + + @_hook.hook + class get_method_hook(Hook[MethodContext, MypyType]): + def extra_init(self) -> None: + self.shared_logic = self.plugin.SharedCallableHookLogic( + fullname=self.fullname, plugin=self.plugin + ) + + def choose(self) -> bool: + return self.shared_logic.choose() + + def run(self, ctx: MethodContext) -> MypyType: + result = self.shared_logic.run(ctx) + if result is not None: + return result + + if self.super_hook is not None: + return self.super_hook(ctx) + + return ctx.default_return_type + + @_hook.hook + class get_function_hook(Hook[FunctionContext, MypyType]): + def extra_init(self) -> None: + self.shared_logic = self.plugin.SharedCallableHookLogic( + fullname=self.fullname, plugin=self.plugin + ) + + def choose(self) -> bool: + return self.shared_logic.choose() + + def run(self, ctx: FunctionContext) -> MypyType: + result = self.shared_logic.run(ctx) if result is not None: return result From ebe6cdf730fc174bbc44b0075f2d2604c154fb67 Mon Sep 17 00:00:00 2001 From: Stephen Moore Date: Tue, 28 May 2024 14:43:58 +1000 Subject: [PATCH 5/5] Make sure there is a useful error on impossible TypeGuard Turns out that mypy plugin system doesn't give us an opportunity to change the typeguard when it relies on a Type Var. This is because we only know the value of the type var where it's called, which is the get_method/function_hook but those can only change the return type and don't have the ability to change the type guard --- docs/api/changelog.rst | 7 ++ extended_mypy_django_plugin/plugin/_plugin.py | 92 ++++++++++++++++++- .../plugin/actions/_type_checker.py | 28 +++++- .../output_builder.py | 20 +++- tests/test_concrete_annotations.py | 29 +++++- tests/test_errors.py | 89 ++++++++++++++++++ 6 files changed, 259 insertions(+), 6 deletions(-) create mode 100644 tests/test_errors.py diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index b35cbb6..e78a66c 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -3,6 +3,13 @@ Changelog --------- +.. _release-0.5.4: + +0.5.4 - TBD + * Will now check return types for methods and functions more thorouhgly + * Will throw errors if a type guard is used with a concrete annotation that uses + a type var (mypy plugin system is limited in a way that makes this impossible to implement) + .. _release-0.5.3: 0.5.3 - 25 May 2024 diff --git a/extended_mypy_django_plugin/plugin/_plugin.py b/extended_mypy_django_plugin/plugin/_plugin.py index e1268a1..b40163d 100644 --- a/extended_mypy_django_plugin/plugin/_plugin.py +++ b/extended_mypy_django_plugin/plugin/_plugin.py @@ -12,11 +12,13 @@ ClassDefContext, DynamicClassDefContext, FunctionContext, + FunctionSigContext, MethodContext, + MethodSigContext, ) from mypy.semanal import SemanticAnalyzer from mypy.typeanal import TypeAnalyser -from mypy.types import CallableType, Instance +from mypy.types import CallableType, FunctionLike, Instance from mypy.types import Type as MypyType from mypy_django_plugin import main from mypy_django_plugin.django.context import DjangoContext @@ -269,6 +271,7 @@ def run(self, ctx: AnalyzeTypeContext) -> MypyType: elif name is Known.DEFAULT_QUERYSET: method = type_analyzer.find_default_queryset + else: assert_never(name) @@ -299,6 +302,9 @@ class SharedCallableHookLogic: """ Shared logic for modifying the return type of methods and functions that use a concrete annotation with a type variable. + + Note that the signature hook will already raise errors if a concrete annotation is + used with a type var in a type guard. """ def __init__(self, fullname: str, plugin: "ExtendedMypyStubs") -> None: @@ -377,3 +383,87 @@ def run(self, ctx: FunctionContext) -> MypyType: return self.super_hook(ctx) return ctx.default_return_type + + class SharedSignatureHookLogic: + """ + Shared logic for modifying the signature of methods and functions. + + This is only used to find cases where a concrete annotation with a type var + is used in a type guard. + + In this situation the mypy plugin system does not provide an opportunity to fully resolve + the type guard. + """ + + def __init__(self, fullname: str, plugin: "ExtendedMypyStubs") -> None: + self.plugin = plugin + self.store = plugin.store + self.fullname = fullname + + def choose(self) -> bool: + """ + Only choose methods and functions that are returning a type guard + """ + if self.fullname.startswith("builtins."): + return False + + sym = self.plugin.lookup_fully_qualified(self.fullname) + if not sym or not sym.node: + return False + + call = getattr(sym.node, "type", None) + if not isinstance(call, CallableType): + return False + + return call.type_guard is not None + + def run(self, ctx: MethodSigContext | FunctionSigContext) -> MypyType | None: + assert isinstance(ctx.api, TypeChecker) + + type_checking = actions.TypeChecking( + self.store, + api=ctx.api, + lookup_info=self.plugin._lookup_info, + ) + + return type_checking.check_typeguard(ctx.context) + + @_hook.hook + class get_function_signature_hook(Hook[FunctionSigContext, FunctionLike]): + def extra_init(self) -> None: + self.shared_logic = self.plugin.SharedSignatureHookLogic( + fullname=self.fullname, plugin=self.plugin + ) + + def choose(self) -> bool: + return self.shared_logic.choose() + + def run(self, ctx: FunctionSigContext) -> FunctionLike: + result = self.shared_logic.run(ctx) + if result is not None: + return ctx.default_signature.copy_modified(ret_type=result) + + if self.super_hook is not None: + return self.super_hook(ctx) + + return ctx.default_signature + + @_hook.hook + class get_method_signature_hook(Hook[MethodSigContext, FunctionLike]): + def extra_init(self) -> None: + self.shared_logic = self.plugin.SharedSignatureHookLogic( + fullname=self.fullname, plugin=self.plugin + ) + + def choose(self) -> bool: + return self.shared_logic.choose() + + def run(self, ctx: MethodSigContext) -> FunctionLike: + result = self.shared_logic.run(ctx) + if result is not None: + return ctx.default_signature.copy_modified(ret_type=result) + + if self.super_hook is not None: + return self.super_hook(ctx) + + return ctx.default_signature diff --git a/extended_mypy_django_plugin/plugin/actions/_type_checker.py b/extended_mypy_django_plugin/plugin/actions/_type_checker.py index dffa596..f03f94d 100644 --- a/extended_mypy_django_plugin/plugin/actions/_type_checker.py +++ b/extended_mypy_django_plugin/plugin/actions/_type_checker.py @@ -62,6 +62,7 @@ def _find_type_vars( @dataclasses.dataclass class BasicTypeInfo: is_type: bool + is_guard: bool api: TypeChecker func: CallableType @@ -72,9 +73,14 @@ class BasicTypeInfo: @classmethod def create(cls, api: TypeChecker, func: CallableType, item: MypyType | None = None) -> Self: is_type: bool = False + is_guard: bool = False if item is None: - item = func.ret_type + if func.type_guard: + is_guard = True + item = func.type_guard + else: + item = func.ret_type item = get_proper_type(item) if isinstance(item, TypeType): @@ -102,6 +108,7 @@ def create(cls, api: TypeChecker, func: CallableType, item: MypyType | None = No func=func, item=item, is_type=is_type, + is_guard=is_guard, type_vars=type_vars, concrete_annotation=concrete_annotation, ) @@ -189,11 +196,30 @@ def _get_info(self, context: Context) -> BasicTypeInfo | None: return BasicTypeInfo.create(api=self.api, func=func) + def check_typeguard(self, context: Context) -> MypyType | None: + info = self._get_info(context) + if info is None: + return None + + if info.is_guard and info.type_vars and info.contains_concrete_annotation: + # Mypy plugin system doesn't currently provide an opportunity to resolve a type guard when it's for a concrete annotation that uses a type var + self.api.fail( + "Can't use a TypeGuard that uses a Concrete Annotation that uses type variables", + context, + ) + return AnyType(TypeOfAny.from_error) + + return None + def modify_return_type(self, ctx: MethodContext | FunctionContext) -> MypyType | None: info = self._get_info(ctx.context) if info is None: return None + if info.is_guard and info.type_vars and info.concrete_annotation is not None: + # Mypy plugin system doesn't currently provide an opportunity to resolve a type guard when it's for a concrete annotation that uses a type var + return None + if not info.contains_concrete_annotation: return None diff --git a/scripts/test_helpers/extended_mypy_django_plugin_test_driver/output_builder.py b/scripts/test_helpers/extended_mypy_django_plugin_test_driver/output_builder.py index 69aeaca..e91fba2 100644 --- a/scripts/test_helpers/extended_mypy_django_plugin_test_driver/output_builder.py +++ b/scripts/test_helpers/extended_mypy_django_plugin_test_driver/output_builder.py @@ -1,7 +1,12 @@ +import importlib.metadata from collections.abc import Iterator from pytest_mypy_plugins import OutputMatcher -from pytest_mypy_plugins.utils import DaemonOutputMatcher, FileOutputMatcher +from pytest_mypy_plugins.utils import ( + DaemonOutputMatcher, + FileOutputMatcher, + extract_output_matchers_from_out, +) from typing_extensions import Self @@ -57,7 +62,17 @@ def daemon_should_not_restart(self) -> Self: def on(self, path: str) -> Self: return self.__class__(build=self._build, target_file=path) + def from_out(self, out: str, regex: bool = False) -> Self: + self._build.result.extend( + extract_output_matchers_from_out( + out, {}, regex=regex, for_daemon=self._build.for_daemon + ) + ) + return self + def add_revealed_type(self, lnum: int, revealed_type: str) -> Self: + if importlib.metadata.version("mypy") == "1.4.0": + revealed_type = revealed_type.replace("type[", "Type[") assert self.target_file is not None self._build.add( self.target_file, lnum, None, "note", f'Revealed type is "{revealed_type}"' @@ -65,6 +80,9 @@ def add_revealed_type(self, lnum: int, revealed_type: str) -> Self: return self def change_revealed_type(self, lnum: int, message: str) -> Self: + if importlib.metadata.version("mypy") == "1.4.0": + message = message.replace("type[", "Type[") + assert self.target_file is not None found: list[FileOutputMatcher] = [] diff --git a/tests/test_concrete_annotations.py b/tests/test_concrete_annotations.py index 2a470ae..7863a9b 100644 --- a/tests/test_concrete_annotations.py +++ b/tests/test_concrete_annotations.py @@ -9,27 +9,50 @@ def _(expected: OutputBuilder) -> None: "main.py", """ from extended_mypy_django_plugin import Concrete, ConcreteQuerySet, DefaultQuerySet + from typing import cast, TypeGuard - from myapp.models import Parent + from myapp.models import Parent, Child1 models: Concrete[Parent] reveal_type(models) qs: ConcreteQuerySet[Parent] reveal_type(qs) + + def check_cls_with_type_guard(cls: type[Parent]) -> TypeGuard[type[Concrete[Parent]]]: + return True + + def check_instance_with_type_guard(cls: Parent) -> TypeGuard[Concrete[Parent]]: + return True + + cls: type[Parent] = Child1 + assert check_cls_with_type_guard(cls) + reveal_type(cls) + + instance: Parent = cast(Child1, None) + assert check_instance_with_type_guard(instance) + reveal_type(instance) """, ) ( expected.on("main.py") .add_revealed_type( - 6, + 7, "Union[myapp.models.Child1, myapp.models.Child2, myapp.models.Child3, myapp2.models.ChildOther]", ) .add_revealed_type( - 9, + 10, "Union[django.db.models.query._QuerySet[myapp.models.Child1, myapp.models.Child1], myapp.models.Child2QuerySet, django.db.models.query._QuerySet[myapp.models.Child3, myapp.models.Child3], django.db.models.query._QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]]", ) + .add_revealed_type( + 20, + "Union[type[myapp.models.Child1], type[myapp.models.Child2], type[myapp.models.Child3], type[myapp2.models.ChildOther]]", + ) + .add_revealed_type( + 24, + "Union[myapp.models.Child1, myapp.models.Child2, myapp.models.Child3, myapp2.models.ChildOther]", + ) ) def test_sees_apps_removed_when_they_still_exist_but_no_longer_installed( diff --git a/tests/test_errors.py b/tests/test_errors.py new file mode 100644 index 0000000..350dd36 --- /dev/null +++ b/tests/test_errors.py @@ -0,0 +1,89 @@ +import importlib.metadata + +from extended_mypy_django_plugin_test_driver import OutputBuilder, Scenario + + +class TestErrors: + def test_cant_use_typevar_concrete_annotation_in_function_or_method_typeguard( + self, scenario: Scenario + ) -> None: + @scenario.run_and_check_mypy_after + def _(expected: OutputBuilder) -> None: + scenario.make_file( + "main.py", + """ + from typing import TypeGuard, TypeVar, cast + + from myapp.models import Child1, Parent + + from extended_mypy_django_plugin import Concrete, ConcreteQuerySet, DefaultQuerySet + + T_Parent = Concrete.type_var("T_Parent", Parent) + + def function_with_type_typeguard( + cls: type[Parent], expect: type[T_Parent] + ) -> TypeGuard[type[Concrete[T_Parent]]]: + return hasattr(cls, "objects") + + cls1: type[Parent] = Child1 + assert function_with_type_typeguard(cls1, Parent) + reveal_type(cls1) + + def function_with_instance_typeguard( + instance: Parent, expect: type[T_Parent] + ) -> TypeGuard[Concrete[T_Parent]]: + return True + + instance1: Parent = cast(Child1, None) + assert function_with_instance_typeguard(instance1, Parent) + reveal_type(instance1) + + class Logic: + def method_with_type_typeguard( + self, cls: type[Parent], expect: type[T_Parent] + ) -> TypeGuard[type[Concrete[T_Parent]]]: + return hasattr(cls, "objects") + + def method_with_instance_typeguard( + self, instance: T_Parent, expect: type[T_Parent] + ) -> TypeGuard[Concrete[T_Parent]]: + return True + + logic = Logic() + cls2: type[Parent] = Child1 + assert logic.method_with_type_typeguard(cls2, Parent) + reveal_type(cls2) + + instance2: Parent = cast(Child1, None) + assert logic.method_with_instance_typeguard(instance2, Parent) + reveal_type(instance2) + """, + ) + + out = """ + main:15: error: Can't use a TypeGuard that uses a Concrete Annotation that uses type variables [misc] + main:15: error: Value of type variable "T_Parent" of "function_with_type_typeguard" cannot be "Parent" [type-var] + main:15: error: Only concrete class can be given where "type[Parent]" is expected [type-abstract] + main:16: note: Revealed type is "type[Concrete?[T_Parent?]]" + main:24: error: Can't use a TypeGuard that uses a Concrete Annotation that uses type variables [misc] + main:24: error: Value of type variable "T_Parent" of "function_with_instance_typeguard" cannot be "Parent" [type-var] + main:24: error: Only concrete class can be given where "type[Parent]" is expected [type-abstract] + main:25: note: Revealed type is "Concrete?[T_Parent?]" + main:40: error: Can't use a TypeGuard that uses a Concrete Annotation that uses type variables [misc] + main:40: error: Value of type variable "T_Parent" of "method_with_type_typeguard" of "Logic" cannot be "Parent" [type-var] + main:40: error: Only concrete class can be given where "type[Parent]" is expected [type-abstract] + main:41: note: Revealed type is "type[Concrete?[T_Parent?]]" + main:44: error: Can't use a TypeGuard that uses a Concrete Annotation that uses type variables [misc] + main:44: error: Value of type variable "T_Parent" of "method_with_instance_typeguard" of "Logic" cannot be "Parent" [type-var] + main:44: error: Only concrete class can be given where "type[Parent]" is expected [type-abstract] + main:45: note: Revealed type is "Concrete?[T_Parent?]" + """ + + if importlib.metadata.version("mypy") == "1.4.0": + out = "\n".join( + line + for line in out.split("\n") + if "Only concrete class can be given" not in line + ).replace('Revealed type is "type', 'Revealed type is "Type') + + expected.from_out(out)