Unify plugin check for model type info (#2263)

Replace `is_model_subclass_info` with `is_model_type` helper
typeddjango · Jul 17, 2024 · ac36393 · ac36393
1 parent b2b1afa
commit ac36393
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 13 deletions.
diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py
@@ -301,10 +301,6 @@ def all_registered_model_classes(self) -> Set[Type[models.Model]]:
 
         return all_model_bases
 
-    @cached_property
-    def all_registered_model_class_fullnames(self) -> Set[str]:
-        return {helpers.get_class_fullname(cls) for cls in self.all_registered_model_classes}
-
     @cached_property
     def model_class_fullnames_by_label(self) -> Mapping[str, str]:
         return {

diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py
@@ -401,12 +401,6 @@ def get_typechecker_api(ctx: Union[AttributeContext, MethodContext, FunctionCont
     return ctx.api
 
 
-def is_model_subclass_info(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 check_types_compatible(
     ctx: Union[FunctionContext, MethodContext], *, expected_type: MypyType, actual_type: MypyType, error_message: str
 ) -> None:

diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py
@@ -174,7 +174,7 @@ def get_function_hook(self, fullname: str) -> Optional[Callable[[FunctionContext
             if info.has_base(fullnames.FIELD_FULLNAME):
                 return partial(fields.transform_into_proper_return_type, django_context=self.django_context)
 
-            if helpers.is_model_subclass_info(info, self.django_context):
+            if helpers.is_model_type(info):
                 return partial(init_create.redefine_and_typecheck_model_init, django_context=self.django_context)
 
         return None

diff --git a/mypy_django_plugin/transformers/fields.py b/mypy_django_plugin/transformers/fields.py
@@ -24,7 +24,7 @@ def _get_current_field_from_assignment(
     ctx: FunctionContext, django_context: DjangoContext
 ) -> Optional[Union["Field[Any, Any]", ForeignObjectRel, "GenericForeignKey"]]:
     outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
-    if outer_model_info is None or not helpers.is_model_subclass_info(outer_model_info, django_context):
+    if outer_model_info is None or not helpers.is_model_type(outer_model_info):
         return None
 
     field_name = None
@@ -234,7 +234,7 @@ def transform_into_proper_return_type(ctx: FunctionContext, django_context: Djan
     assert isinstance(default_return_type, Instance)
 
     outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
-    if outer_model_info is None or not helpers.is_model_subclass_info(outer_model_info, django_context):
+    if outer_model_info is None or not helpers.is_model_type(outer_model_info):
         return ctx.default_return_type
 
     assert isinstance(outer_model_info, TypeInfo)