Fix .filter-kwargs lookup crash during cached runs (#828)

* Fix .filter-kwargs lookup crash during cached runs

* Update mypy_django_plugin/lib/helpers.py

Co-authored-by: Nikita Sobolev <[email protected]>

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

mypy_django_plugin/transformers/managers.py

@@ -204,21 +204,21 @@ 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.
             #
             # 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:

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]
 

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()