diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index c417d82..d6d7f9c 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -10,8 +10,8 @@ from mypy import checker from mypy.checker import TypeChecker from mypy.mro import calculate_mro from mypy.nodes import ( - GDEF, MDEF, Argument, Block, ClassDef, Expression, FuncDef, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode, - SymbolTable, SymbolTableNode, TypeInfo, Var, + GDEF, MDEF, Argument, Block, ClassDef, Expression, FuncDef, MemberExpr, MypyFile, NameExpr, PlaceholderNode, + StrExpr, SymbolNode, SymbolTable, SymbolTableNode, TypeInfo, Var, ) from mypy.plugin import ( AttributeContext, CheckerPluginInterface, ClassDefContext, DynamicClassDefContext, FunctionContext, MethodContext, @@ -309,39 +309,67 @@ def add_new_sym_for_info(info: TypeInfo, *, name: str, sym_type: MypyType) -> No plugin_generated=True) -def _prepare_new_method_arguments(node: FuncDef) -> Tuple[List[Argument], MypyType]: - arguments = [] - for argument in node.arguments[1:]: - if argument.type_annotation is None: - argument.type_annotation = AnyType(TypeOfAny.unannotated) - arguments.append(argument) - - if isinstance(node.type, CallableType): - return_type = node.type.ret_type - else: - return_type = AnyType(TypeOfAny.unannotated) - - return arguments, return_type +def build_unannotated_method_args(method_node: FuncDef) -> Tuple[List[Argument], MypyType]: + prepared_arguments = [] + for argument in method_node.arguments[1:]: + argument.type_annotation = AnyType(TypeOfAny.unannotated) + prepared_arguments.append(argument) + return_type = AnyType(TypeOfAny.unannotated) + return prepared_arguments, return_type def copy_method_to_another_class(ctx: ClassDefContext, self_type: Instance, new_method_name: str, method_node: FuncDef) -> None: - arguments, return_type = _prepare_new_method_arguments(method_node) - semanal_api = get_semanal_api(ctx) - for argument in arguments: - if argument.type_annotation is not None: - argument.type_annotation = semanal_api.anal_type(argument.type_annotation, - allow_placeholder=True) + if method_node.type is None: + if not semanal_api.final_iteration: + semanal_api.defer() + return - if return_type is not None: - ret = semanal_api.anal_type(return_type, - allow_placeholder=True) - assert ret is not None - return_type = ret + arguments, return_type = build_unannotated_method_args(method_node) + add_method(ctx, + new_method_name, + args=arguments, + return_type=return_type, + self_type=self_type) + return + + method_type = method_node.type + if not isinstance(method_type, CallableType): + if not semanal_api.final_iteration: + semanal_api.defer() + return + + arguments = [] + bound_return_type = semanal_api.anal_type(method_type.ret_type, + allow_placeholder=True) + assert bound_return_type is not None + + if isinstance(bound_return_type, PlaceholderNode): + return + + for arg_name, arg_type, original_argument in zip(method_type.arg_names[1:], + method_type.arg_types[1:], + method_node.arguments[1:]): + bound_arg_type = semanal_api.anal_type(arg_type, allow_placeholder=True) + assert bound_arg_type is not None + + if isinstance(bound_arg_type, PlaceholderNode): + return + + var = Var(name=original_argument.variable.name, + type=arg_type) + var.line = original_argument.variable.line + var.column = original_argument.variable.column + argument = Argument(variable=var, + type_annotation=bound_arg_type, + initializer=original_argument.initializer, + kind=original_argument.kind) + argument.set_line(original_argument) + arguments.append(argument) add_method(ctx, new_method_name, args=arguments, - return_type=return_type, + return_type=bound_return_type, self_type=self_type) diff --git a/test-data/typecheck/managers/querysets/test_from_queryset.yml b/test-data/typecheck/managers/querysets/test_from_queryset.yml index 1562844..e9f2ad4 100644 --- a/test-data/typecheck/managers/querysets/test_from_queryset.yml +++ b/test-data/typecheck/managers/querysets/test_from_queryset.yml @@ -146,3 +146,36 @@ return param NewManager = models.Manager.from_queryset(ModelQuerySet) + +- case: from_queryset_with_inherited_manager_and_typing_no_return + disable_cache: true + main: | + from myapp.models import MyModel + reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]' + reveal_type(MyModel().objects.get()) # N: Revealed type is 'myapp.models.MyModel*' + reveal_type(MyModel().objects.base_queryset_method) # N: Revealed type is 'def (param: Union[builtins.int, builtins.str]) -> ' + reveal_type(MyModel().objects.base_queryset_method(2)) # N: Revealed type is '' + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from myapp.managers import NewManager + class MyModel(models.Model): + objects = NewManager() + - path: myapp/managers.py + content: | + from django.db import models + from myapp.base_queryset import BaseQuerySet + class ModelQuerySet(BaseQuerySet): + pass + NewManager = models.Manager.from_queryset(ModelQuerySet) + - path: myapp/base_queryset.py + content: | + from typing import NoReturn, Union + from django.db import models + class BaseQuerySet(models.QuerySet): + def base_queryset_method(self, param: Union[int, str]) -> NoReturn: + raise ValueError \ No newline at end of file diff --git a/test-data/typecheck/managers/test_managers.yml b/test-data/typecheck/managers/test_managers.yml index 22dbf33..82206d6 100644 --- a/test-data/typecheck/managers/test_managers.yml +++ b/test-data/typecheck/managers/test_managers.yml @@ -307,15 +307,15 @@ - case: custom_manager_returns_proper_model_types main: | from myapp.models import User - reveal_type(User.objects) # N: Revealed type is 'myapp.models.User_MyManager[myapp.models.User]' - reveal_type(User.objects.select_related()) # N: Revealed type is 'myapp.models.User_MyManager[myapp.models.User]' + reveal_type(User.objects) # N: Revealed type is 'myapp.models.User_MyManager2[myapp.models.User]' + reveal_type(User.objects.select_related()) # N: Revealed type is 'myapp.models.User_MyManager2[myapp.models.User]' reveal_type(User.objects.get()) # N: Revealed type is 'myapp.models.User*' reveal_type(User.objects.get_instance()) # N: Revealed type is 'builtins.int' reveal_type(User.objects.get_instance_untyped('hello')) # N: Revealed type is 'Any' from myapp.models import ChildUser - reveal_type(ChildUser.objects) # N: Revealed type is 'myapp.models.ChildUser_MyManager[myapp.models.ChildUser]' - reveal_type(ChildUser.objects.select_related()) # N: Revealed type is 'myapp.models.ChildUser_MyManager[myapp.models.ChildUser]' + reveal_type(ChildUser.objects) # N: Revealed type is 'myapp.models.ChildUser_MyManager2[myapp.models.ChildUser]' + reveal_type(ChildUser.objects.select_related()) # N: Revealed type is 'myapp.models.ChildUser_MyManager2[myapp.models.ChildUser]' reveal_type(ChildUser.objects.get()) # N: Revealed type is 'myapp.models.ChildUser*' reveal_type(ChildUser.objects.get_instance()) # N: Revealed type is 'builtins.int' reveal_type(ChildUser.objects.get_instance_untyped('hello')) # N: Revealed type is 'Any'