add support for Apps.get_model for migrations

mypy_django_plugin/plugins/migrations.py

@@ -0,0 +1,24 @@
+from typing import cast
+
+from mypy.checker import TypeChecker
+from mypy.nodes import TypeInfo
+from mypy.plugin import MethodContext
+from mypy.types import Type, Instance, TypeType
+
+from mypy_django_plugin import helpers
+
+
+def determine_model_cls_from_string_for_migrations(ctx: MethodContext) -> Type:
+    app_label = ctx.args[ctx.callee_arg_names.index('app_label')][0].value
+    model_name = ctx.args[ctx.callee_arg_names.index('model_name')][0].value
+
+    api = cast(TypeChecker, ctx.api)
+    model_fullname = helpers.get_model_fullname(app_label, model_name, all_modules=api.modules)
+
+    if model_fullname is None:
+        return ctx.default_return_type
+    model_info = helpers.lookup_fully_qualified_generic(model_fullname,
+                                                        all_modules=api.modules)
+    if model_info is None or not isinstance(model_info, TypeInfo):
+        return ctx.default_return_type
+    return TypeType(Instance(model_info, []))

mypy_django_plugin/plugins/models.py

@@ -3,7 +3,7 @@ from abc import abstractmethod, ABCMeta
 from typing import cast, Iterator, Tuple, Optional, Dict
 
 from mypy.nodes import ClassDef, AssignmentStmt, CallExpr, MemberExpr, StrExpr, NameExpr, MDEF, TypeInfo, Var, SymbolTableNode, \
-    Lvalue, Expression, MypyFile
+    Lvalue, Expression, MypyFile, Context
 from mypy.plugin import ClassDefContext
 from mypy.semanal import SemanticAnalyzerPass2
 from mypy.types import Instance
@@ -128,9 +128,15 @@ class AddRelatedManagers(ModelClassInitializer):
             for defn in iter_over_classdefs(module_file):
                 for lvalue, rvalue in iter_call_assignments(defn):
                     if is_related_field(rvalue, module_file):
-                        ref_to_fullname = extract_ref_to_fullname(rvalue,
-                                                                  module_file=module_file,
-                                                                  all_modules=self.api.modules)
+                        try:
+                            ref_to_fullname = extract_ref_to_fullname(rvalue,
+                                                                      module_file=module_file,
+                                                                      all_modules=self.api.modules)
+                        except helpers.InvalidModelString as exc:
+                            self.api.fail(f'Invalid value for a to= parameter: {exc.model_string!r}',
+                                          Context(line=rvalue.line))
+                            return None
+
                         if self.model_classdef.fullname == ref_to_fullname:
                             if 'related_name' in rvalue.arg_names:
                                 related_name_expr = rvalue.args[rvalue.arg_names.index('related_name')]

mypy_django_plugin/plugins/related_fields.py

@@ -2,7 +2,7 @@ import typing
 from typing import Optional, cast
 
 from mypy.checker import TypeChecker
-from mypy.nodes import StrExpr, TypeInfo
+from mypy.nodes import StrExpr, TypeInfo, Context
 from mypy.plugin import FunctionContext
 from mypy.types import Type, CallableType, Instance, AnyType, TypeOfAny
 
@@ -57,7 +57,12 @@ def get_valid_to_value_or_none(ctx: FunctionContext) -> Optional[Instance]:
 
 
 def extract_to_parameter_as_get_ret_type_for_related_field(ctx: FunctionContext) -> Type:
-    referred_to_type = get_valid_to_value_or_none(ctx)
+    try:
+        referred_to_type = get_valid_to_value_or_none(ctx)
+    except helpers.InvalidModelString as exc:
+        ctx.api.fail(f'Invalid value for a to= parameter: {exc.model_string!r}', ctx.context)
+        return fill_typevars_with_any(ctx.default_return_type)
+
     if referred_to_type is None:
         # couldn't extract to= value
         return fill_typevars_with_any(ctx.default_return_type)