replace field callback with new API

Author: mkurnikov

mypy_django_plugin/django/context.py

@@ -16,6 +16,7 @@
 from django.db.models.sql.query import Query
 from django.utils.functional import cached_property
 from mypy.checker import TypeChecker
+from mypy.nodes import TypeInfo
 from mypy.plugin import MethodContext
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
@@ -214,6 +215,10 @@ def all_registered_model_classes(self) -> Set[Type[models.Model]]:
     def all_registered_model_class_fullnames(self) -> Set[str]:
         return {helpers.get_class_fullname(cls) for cls in self.all_registered_model_classes}
 
+    def is_model_subclass(self, info: TypeInfo) -> bool:
+        return (info.fullname in self.all_registered_model_class_fullnames
+                or info.has_base(fullnames.MODEL_CLASS_FULLNAME))
+
     def get_attname(self, field: Field) -> str:
         attname = field.attname
         return attname

mypy_django_plugin/lib/helpers.py

@@ -20,7 +20,6 @@
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny, UnionType
 
-from mypy_django_plugin.lib import fullnames
 from mypy_django_plugin.transformers2 import new_helpers
 
 if TYPE_CHECKING:
@@ -201,10 +200,12 @@ def get_method_return_type(self) -> MypyType:
 
 class GetFunctionCallback(TypeCheckerPluginCallback):
     ctx: FunctionContext
+    default_return_type: MypyType
 
     def __call__(self, ctx: FunctionContext) -> MypyType:
         self.type_checker = cast(TypeChecker, ctx.api)
         self.ctx = ctx
+        self.default_return_type = ctx.default_return_type
         return self.get_function_return_type()
 
     @abstractmethod
@@ -403,11 +404,6 @@ def resolve_string_attribute_value(attr_expr: Expression, django_context: 'Djang
     return None
 
 
-def is_subclass_of_model(info: TypeInfo, django_context: 'DjangoContext') -> bool:
-    return (info.fullname in django_context.all_registered_model_class_fullnames
-            or info.has_base(fullnames.MODEL_CLASS_FULLNAME))
-
-
 def new_typeinfo(name: str,
                  *,
                  bases: List[Instance],

mypy_django_plugin/main.py

@@ -14,9 +14,8 @@
 import mypy_django_plugin.transformers.orm_lookups
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
-from mypy_django_plugin.transformers import (
-    fields, init_create, meta, querysets,
-)
+from mypy_django_plugin.transformers import init_create, meta, querysets
+from mypy_django_plugin.transformers2.fields import FieldContructorCallback
 from mypy_django_plugin.transformers2.forms import (
     FormCallback, GetFormCallback, GetFormClassCallback,
 )
@@ -167,9 +166,9 @@ def get_function_hook(self, fullname: str
         info = self._get_typeinfo_or_none(fullname)
         if info:
             if info.has_base(fullnames.FIELD_FULLNAME):
-                return partial(fields.transform_into_proper_return_type, django_context=self.django_context)
+                return FieldContructorCallback(self)
 
-            if helpers.is_subclass_of_model(info, self.django_context):
+            if self.django_context.is_model_subclass(info):
                 return partial(init_create.redefine_and_typecheck_model_init, django_context=self.django_context)
         return None