try and use named arguments from caller for matching name (#24)

Co-authored-by: Anthony Sottile <[email protected]>

Author: asottile-sentry

mypy_django_plugin/lib/helpers.py

@@ -10,6 +10,7 @@
 from mypy.nodes import (
     GDEF,
     MDEF,
+    ArgKind,
     AssignmentStmt,
     Block,
     ClassDef,
@@ -175,6 +176,12 @@ def get_call_argument_by_name(ctx: Union[FunctionContext, MethodContext], name:
     Return the expression for the specific argument.
     This helper should only be used with non-star arguments.
     """
+    # try and pull the named argument from the caller first
+    for kinds, argnames, args in zip(ctx.arg_kinds, ctx.arg_names, ctx.args):
+        for kind, argname, arg in zip(kinds, argnames, args):
+            if kind == ArgKind.ARG_NAMED and argname == name:
+                return arg
+
     if name not in ctx.callee_arg_names:
         return None
     idx = ctx.callee_arg_names.index(name)

tests/typecheck/fields/test_related.yml

@@ -800,6 +800,31 @@
                 class User(AbstractUser):
                     pass
 
+-   case: nullable_foreign_key_with_init_overridden
+    main: |
+        from myapp.models import A
+        reveal_type(A.objects.get().b)  # N: Revealed type is "Union[myapp.models.B, None]"
+    installed_apps:
+        -   myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from typing import Any, TypeVar
+                from django.db import models
+
+                _ST = TypeVar("_ST", contravariant=True)
+                _GT = TypeVar("_GT", covariant=True)
+
+                class FK(models.ForeignKey[_ST, _GT]):
+                    def __init__(self, *args: Any, **kwargs: Any) -> None:
+                        kwargs.setdefault('on_delete', models.CASCADE)
+                        super().__init__(*args, **kwargs)
+
+                class B(models.Model): ...
+
+                class A(models.Model):
+                    b = FK(B, null=True, on_delete=models.CASCADE)
 
 -   case: related_manager_is_a_subclass_of_default_manager
     main: |