From 5040599486e27f03a3d865293d4340e987f1a803 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Tue, 16 Jul 2024 22:12:00 +0200 Subject: [PATCH] Unify plugin check for model type info Replace `is_model_subclass_info` with `is_model_type` helper --- mypy_django_plugin/django/context.py | 4 ---- mypy_django_plugin/lib/helpers.py | 6 ------ mypy_django_plugin/main.py | 2 +- mypy_django_plugin/transformers/fields.py | 4 ++-- 4 files changed, 3 insertions(+), 13 deletions(-) diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index 3a6af7042..146ca368e 100644 --- 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 index de8aeb1b4..3332ae9b2 100644 --- 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 index e067d4987..a2f52ab2e 100644 --- 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 index cfd280aa4..9de0b726f 100644 --- 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)