diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index d4ddfb134..9edfc7de6 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -406,30 +406,31 @@ def copy_method_to_another_class( return_type = bind_or_analyze_type(method_type.ret_type, semanal_api, original_module_name) if return_type is None: return - try: - original_arguments = method_node.arguments[1:] - except AttributeError: - original_arguments = [] + # We build the arguments from the method signature (`CallableType`), because if we were to + # use the arguments from the method node (`FuncDef.arguments`) we're not compatible with + # a method loaded from cache. As mypy doesn't serialize `FuncDef.arguments` when caching arguments = [] - for arg_name, arg_type, original_argument in zip( - method_type.arg_names[1:], method_type.arg_types[1:], original_arguments + # Note that the first argument is excluded, as that's `self` + for pos, (arg_type, arg_kind, arg_name) in enumerate( + zip(method_type.arg_types[1:], method_type.arg_kinds[1:], method_type.arg_names[1:]), + start=1, ): bound_arg_type = bind_or_analyze_type(arg_type, semanal_api, original_module_name) if bound_arg_type is None: 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, + if arg_name is None and hasattr(method_node, "arguments"): + arg_name = method_node.arguments[pos].variable.name + arguments.append( + Argument( + # Positional only arguments can have name as `None`, if we can't find a name, we just invent one.. + variable=Var(name=arg_name if arg_name is not None else str(pos), type=arg_type), + type_annotation=bound_arg_type, + initializer=None, + kind=arg_kind, + pos_only=arg_name is None, + ) ) - argument.set_line(original_argument) - arguments.append(argument) add_method_to_class( semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index c86886ee1..09c2ea6c6 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -204,13 +204,8 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME: break for name, sym in class_mro_info.names.items(): - if isinstance(sym.node, FuncDef): - func_node = sym.node - elif isinstance(sym.node, Decorator): - func_node = sym.node.func - else: + if not isinstance(sym.node, (FuncDef, Decorator)): continue - # Insert the queryset method name as a class member. Note that the type of # the method is set as Any. Figuring out the type is the job of the # 'resolve_manager_method' attribute hook, which comes later. @@ -218,7 +213,12 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte # class BaseManagerFromMyQuerySet(BaseManager): # queryset_method: Any = ... # - helpers.add_new_sym_for_info(new_manager_info, name=name, sym_type=AnyType(TypeOfAny.special_form)) + helpers.add_new_sym_for_info( + new_manager_info, + name=name, + sym_type=AnyType(TypeOfAny.special_form), + no_serialize=True, + ) # we need to copy all methods in MRO before django.db.models.query.QuerySet # Gather names of all BaseManager methods @@ -278,7 +278,9 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte ) # Insert the new manager (dynamic) class - assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True)) + assert semanal_api.add_symbol_table_node( + ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True, no_serialize=True) + ) def fail_if_manager_type_created_in_model_body(ctx: MethodContext) -> MypyType: diff --git a/mypy_django_plugin/transformers/orm_lookups.py b/mypy_django_plugin/transformers/orm_lookups.py index e3adc1d6b..67fbd1ce0 100644 --- a/mypy_django_plugin/transformers/orm_lookups.py +++ b/mypy_django_plugin/transformers/orm_lookups.py @@ -9,6 +9,9 @@ def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext) -> MypyType: + # Expected formal arguments for filter methods are `*args` and `**kwargs`. We'll only typecheck + # `**kwargs`, which means that `arg_names[1]` is what we're interested in. + lookup_kwargs = ctx.arg_names[1] provided_lookup_types = ctx.arg_types[1] diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index a9c9cecd5..07f921569 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -346,3 +346,41 @@ base_manager = BaseManagerFromMyQuerySet() manager = ManagerFromMyQuerySet() custom_manager = MyManagerFromMyQuerySet() + +- case: from_queryset_includes_methods_returning_queryset + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects.none) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.exclude) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.complex_filter) # N: Revealed type is "def (filter_obj: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.union) # N: Revealed type is "def (*other_qs: Any, *, all: builtins.bool =) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.intersection) # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.difference) # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.select_for_update) # N: Revealed type is "def (nowait: builtins.bool =, skip_locked: builtins.bool =, of: typing.Sequence[builtins.str] =, no_key: builtins.bool =) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.select_related) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.prefetch_related) # N: Revealed type is "def (*lookups: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.annotate) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.alias) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.order_by) # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.distinct) # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.reverse) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.defer) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.only) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + reveal_type(MyModel.objects.using) # N: Revealed type is "def (alias: Union[builtins.str, None]) -> myapp.models.MyQuerySet[myapp.models.MyModel*]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class MyQuerySet(models.QuerySet["MyModel"]): + ... + + MyManager = BaseManager.from_queryset(MyQuerySet) + class MyModel(models.Model): + objects = MyManager()