diff --git a/django-stubs/apps/registry.pyi b/django-stubs/apps/registry.pyi index 69381ad..15464f1 100644 --- a/django-stubs/apps/registry.pyi +++ b/django-stubs/apps/registry.pyi @@ -24,7 +24,7 @@ class Apps: def get_app_configs(self) -> Iterable[AppConfig]: ... def get_app_config(self, app_label: str) -> AppConfig: ... # it's not possible to support it in plugin properly now - def get_models(self, include_auto_created: bool = ..., include_swapped: bool = ...) -> List[Type[Any]]: ... + def get_models(self, include_auto_created: bool = ..., include_swapped: bool = ...) -> List[Type[Model]]: ... def get_model(self, app_label: str, model_name: Optional[str] = ..., require_ready: bool = ...) -> Type[Any]: ... def register_model(self, app_label: str, model: Type[Model]) -> None: ... def is_installed(self, app_name: str) -> bool: ... diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index 0a29653..2cdd3b2 100644 --- a/mypy_django_plugin/django/context.py +++ b/mypy_django_plugin/django/context.py @@ -1,20 +1,18 @@ import os from collections import defaultdict from contextlib import contextmanager -from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Type +from typing import Dict, Iterator, Optional, Set, TYPE_CHECKING, Tuple, Type from django.core.exceptions import FieldError from django.db.models.base import Model -from django.db.models.fields import AutoField, CharField, Field from django.db.models.fields.related import ForeignKey, RelatedField from django.db.models.fields.reverse_related import ForeignObjectRel from django.db.models.sql.query import Query from django.utils.functional import cached_property from mypy.checker import TypeChecker -from mypy.types import AnyType, Instance -from mypy.types import Type as MypyType -from mypy.types import TypeOfAny +from mypy.types import AnyType, Instance, Type as MypyType, TypeOfAny +from django.db.models.fields import AutoField, CharField, Field from mypy_django_plugin.lib import helpers try: @@ -170,20 +168,26 @@ class DjangoContext: self.settings = settings @cached_property - def model_modules(self) -> Dict[str, List[Type[Model]]]: + def model_modules(self) -> Dict[str, Set[Type[Model]]]: """ All modules that contain Django models. """ if self.apps_registry is None: return {} - modules: Dict[str, List[Type[Model]]] = defaultdict(list) - for model_cls in self.apps_registry.get_models(): - modules[model_cls.__module__].append(model_cls) + modules: Dict[str, Set[Type[Model]]] = defaultdict(set) + for concrete_model_cls in self.apps_registry.get_models(): + modules[concrete_model_cls.__module__].add(concrete_model_cls) + # collect abstract=True models + for model_cls in concrete_model_cls.mro()[1:]: + if (issubclass(model_cls, Model) + and hasattr(model_cls, '_meta') + and model_cls._meta.abstract): + modules[model_cls.__module__].add(model_cls) return modules def get_model_class_by_fullname(self, fullname: str) -> Optional[Type[Model]]: # Returns None if Model is abstract module, _, model_cls_name = fullname.rpartition('.') - for model_cls in self.model_modules.get(module, []): + for model_cls in self.model_modules.get(module, set()): if model_cls.__name__ == model_cls_name: return model_cls return None diff --git a/test-data/typecheck/fields/test_related.yml b/test-data/typecheck/fields/test_related.yml index c310277..1c8dd60 100644 --- a/test-data/typecheck/fields/test_related.yml +++ b/test-data/typecheck/fields/test_related.yml @@ -544,3 +544,31 @@ class Author(models.Model): blogs = models.ManyToManyField(Blog) file = models.FileField() + + +- case: test_foreign_key_from_superclass_inherits_correctly + main: | + from myapp.models import MyUser, Book, Article, LibraryEntity + reveal_type(Book().registered_by_user) # N: Revealed type is 'myapp.models.MyUser*' + reveal_type(Article().registered_by_user) # N: Revealed type is 'myapp.models.MyUser*' + + user = MyUser() + reveal_type(user.book_set) # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp.models.Book]' + reveal_type(user.article_set) # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp.models.Article]' + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + class MyUser(models.Model): + pass + class LibraryEntity(models.Model): + class Meta: + abstract = True + registered_by_user = models.ForeignKey(MyUser, on_delete=models.CASCADE) + class Book(LibraryEntity): + pass + class Article(LibraryEntity): + pass \ No newline at end of file