Plugin misc improvements

mypy_django_plugin/lib/helpers.py

@@ -202,7 +202,12 @@ def get_min_argument_count(ctx: MethodContext | FunctionContext) -> int:
 def get_call_argument_by_name(ctx: FunctionContext | MethodContext, name: str) -> Expression | None:
     """
     Return the expression for the specific argument.
-    This helper should only be used with non-star arguments.
+    This helper supports named and positional arguments and should only be used with non-star arguments.
+
+    Ex:
+        Given `def my_func(a: int, b: str), `get_call_argument_by_name(ctx, "b")` works for`
+        - `my_func(1, b="x")`
+        - `my_func(1, "x")`
     """
     # try and pull the named argument from the caller first
     for kinds, argnames, args in zip(ctx.arg_kinds, ctx.arg_names, ctx.args, strict=False):
@@ -212,6 +217,7 @@ def get_call_argument_by_name(ctx: FunctionContext | MethodContext, name: str) -
 
     if name not in ctx.callee_arg_names:
         return None
+
     idx = ctx.callee_arg_names.index(name)
     args = ctx.args[idx]
     if len(args) != 1:

mypy_django_plugin/transformers/init_create.py

@@ -59,40 +59,34 @@ def typecheck_model_method(
 
 def typecheck_model_init(ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
     default_return_type = get_proper_type(ctx.default_return_type)
-    assert isinstance(default_return_type, Instance)
-
-    model_cls = django_context.get_model_class_by_fullname(default_return_type.type.fullname)
-    if model_cls is not None:
+    if (
+        isinstance(default_return_type, Instance)
+        and (model_cls := django_context.get_model_class_by_fullname(default_return_type.type.fullname)) is not None
+    ):
         typecheck_model_method(ctx, django_context, model_cls, "__init__")
 
     return ctx.default_return_type
 
 
 def typecheck_model_create(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
     default_return_type = get_proper_type(ctx.default_return_type)
-    if not isinstance(default_return_type, Instance):
-        # only work with ctx.default_return_type = model Instance
-        return ctx.default_return_type
-
-    model_cls = django_context.get_model_class_by_fullname(default_return_type.type.fullname)
-    if model_cls is not None:
+    if (
+        isinstance(default_return_type, Instance)
+        and (model_cls := django_context.get_model_class_by_fullname(default_return_type.type.fullname)) is not None
+    ):
         typecheck_model_method(ctx, django_context, model_cls, "create")
 
     return ctx.default_return_type
 
 
 def typecheck_model_acreate(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
     default_return_type = get_proper_type(ctx.default_return_type)
-    if not isinstance(default_return_type, Instance):
-        return ctx.default_return_type
-
-    # default_return_type at this point should be of type Coroutine[Any, Any, <Model>]
-    model = get_proper_type(default_return_type.args[-1])
-    if not isinstance(model, Instance):
-        return ctx.default_return_type
-
-    model_cls = django_context.get_model_class_by_fullname(model.type.fullname)
-    if model_cls is not None:
+    if (
+        isinstance(default_return_type, Instance)
+        # default_return_type at this point should be of type Coroutine[Any, Any, <Model>]
+        and isinstance((model := get_proper_type(default_return_type.args[-1])), Instance)
+        and (model_cls := django_context.get_model_class_by_fullname(model.type.fullname)) is not None
+    ):
         typecheck_model_method(ctx, django_context, model_cls, "acreate")
 
     return ctx.default_return_type

mypy_django_plugin/transformers/manytomany.py

@@ -32,18 +32,15 @@ def fill_model_args_for_many_to_many_field(
     django_context: DjangoContext,
 ) -> MypyType:
     default_return_type = get_proper_type(ctx.default_return_type)
-    if (
-        not ctx.args
-        or not ctx.args[0]
-        or not isinstance(default_return_type, Instance)
-        or len(default_return_type.args) < 2
+    if not (
+        ctx.args
+        and ctx.args[0]
+        and isinstance(default_return_type, Instance)
+        and len(default_return_type.args) >= 2
+        and (args := get_m2m_arguments(ctx=ctx, model_info=model_info, django_context=django_context))
     ):
         return ctx.default_return_type
 
-    args = get_m2m_arguments(ctx=ctx, model_info=model_info, django_context=django_context)
-    if args is None:
-        return ctx.default_return_type
-
     default_to_arg = get_proper_type(default_return_type.args[0])
     to_arg: MypyType
     if isinstance(default_to_arg, UninhabitedType):

mypy_django_plugin/transformers/meta.py

@@ -7,23 +7,13 @@
 
 
 def return_proper_field_type_from_get_field(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
-    # Options instance
-    assert isinstance(ctx.type, Instance)
-
-    # bail if list of generic params is empty
-    if len(ctx.type.args) == 0:
-        return ctx.default_return_type
-
-    model_type = get_proper_type(ctx.type.args[0])
-    if not isinstance(model_type, Instance):
-        return ctx.default_return_type
-
-    field_name_expr = helpers.get_call_argument_by_name(ctx, "field_name")
-    if field_name_expr is None:
-        return ctx.default_return_type
-
-    field_name = helpers.resolve_string_attribute_value(field_name_expr, django_context)
-    if field_name is None:
+    if not (
+        isinstance(ctx.type, Instance)
+        and ctx.type.args
+        and isinstance(model_type := get_proper_type(ctx.type.args[0]), Instance)
+        and (field_name_expr := helpers.get_call_argument_by_name(ctx, "field_name")) is not None
+        and (field_name := helpers.resolve_string_attribute_value(field_name_expr, django_context)) is not None
+    ):
         return ctx.default_return_type
 
     field_type = get_field_type_from_model_type_info(model_type.type, field_name)