diff --git a/mypy_django_plugin_newsemanal/django/context.py b/mypy_django_plugin_newsemanal/django/context.py index 3367f56..bf67cda 100644 --- a/mypy_django_plugin_newsemanal/django/context.py +++ b/mypy_django_plugin_newsemanal/django/context.py @@ -12,6 +12,7 @@ from pytest_mypy.utils import temp_environ from django.contrib.postgres.fields import ArrayField from django.db.models.fields import CharField, Field +from django.db.models.fields.reverse_related import ForeignObjectRel from mypy_django_plugin_newsemanal.lib import helpers if TYPE_CHECKING: @@ -122,6 +123,11 @@ class DjangoContext: if isinstance(field, Field): yield field + def get_model_relations(self, model_cls: Type[Model]) -> Iterator[ForeignObjectRel]: + for field in model_cls._meta.get_fields(): + if isinstance(field, ForeignObjectRel): + yield field + def get_primary_key_field(self, model_cls: Type[Model]) -> Field: for field in model_cls._meta.get_fields(): if isinstance(field, Field): diff --git a/mypy_django_plugin_newsemanal/main.py b/mypy_django_plugin_newsemanal/main.py index ea8b2fc..2205047 100644 --- a/mypy_django_plugin_newsemanal/main.py +++ b/mypy_django_plugin_newsemanal/main.py @@ -8,6 +8,7 @@ from mypy.options import Options from mypy.plugin import ClassDefContext, FunctionContext, Plugin, MethodContext from mypy.types import Type as MypyType +from django.db.models.fields.related import RelatedField from mypy_django_plugin_newsemanal.django.context import DjangoContext from mypy_django_plugin_newsemanal.lib import fullnames, metadata from mypy_django_plugin_newsemanal.transformers import fields, settings, querysets, init_create @@ -96,8 +97,15 @@ class NewSemanalDjangoPlugin(Plugin): return [] deps = set() for model_class in defined_model_classes: - for related_object in model_class._meta.related_objects: - related_model_module = related_object.related_model.__module__ + # forward relations + for field in self.django_context.get_model_fields(model_class): + if isinstance(field, RelatedField): + related_model_module = field.related_model.__module__ + if related_model_module != file.fullname(): + deps.add(self._new_dependency(related_model_module)) + # reverse relations + for relation in model_class._meta.related_objects: + related_model_module = relation.related_model.__module__ if related_model_module != file.fullname(): deps.add(self._new_dependency(related_model_module)) return list(deps) @@ -114,7 +122,7 @@ class NewSemanalDjangoPlugin(Plugin): info = self._get_typeinfo_or_none(fullname) if info: if info.has_base(fullnames.FIELD_FULLNAME): - return partial(fields.process_field_instantiation, django_context=self.django_context) + return partial(fields.transform_into_proper_return_type, django_context=self.django_context) if info.has_base(fullnames.MODEL_CLASS_FULLNAME): return partial(init_create.redefine_and_typecheck_model_init, django_context=self.django_context) diff --git a/mypy_django_plugin_newsemanal/transformers/fields.py b/mypy_django_plugin_newsemanal/transformers/fields.py index ad17bba..6c5a390 100644 --- a/mypy_django_plugin_newsemanal/transformers/fields.py +++ b/mypy_django_plugin_newsemanal/transformers/fields.py @@ -105,7 +105,7 @@ def get_referred_to_model_fullname(ctx: FunctionContext, django_context: DjangoC def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType: referred_to_fullname = get_referred_to_model_fullname(ctx, django_context) referred_to_typeinfo = helpers.lookup_fully_qualified_generic(referred_to_fullname, ctx.api.modules) - assert isinstance(referred_to_typeinfo, TypeInfo) + assert isinstance(referred_to_typeinfo, TypeInfo), f'Cannot resolve {referred_to_fullname!r}' referred_to_type = Instance(referred_to_typeinfo, []) default_related_field_type = set_descriptor_types_for_field(ctx) @@ -145,12 +145,6 @@ def transform_into_proper_return_type(ctx: FunctionContext, django_context: Djan return set_descriptor_types_for_field(ctx) -def process_field_instantiation(ctx: FunctionContext, django_context: DjangoContext) -> MypyType: - # Parse __init__ parameters of field into corresponding Model's metadata - # parse_field_init_arguments_into_model_metadata(ctx) - return transform_into_proper_return_type(ctx, django_context) - - def determine_type_of_array_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType: default_return_type = set_descriptor_types_for_field(ctx) diff --git a/mypy_django_plugin_newsemanal/transformers/models.py b/mypy_django_plugin_newsemanal/transformers/models.py index cb92afc..a1b02f7 100644 --- a/mypy_django_plugin_newsemanal/transformers/models.py +++ b/mypy_django_plugin_newsemanal/transformers/models.py @@ -9,8 +9,9 @@ from mypy.plugin import ClassDefContext from mypy.types import Instance from django.db.models.fields import Field +from django.db.models.fields.reverse_related import ManyToOneRel, OneToOneRel, ManyToManyRel from mypy_django_plugin_newsemanal.django.context import DjangoContext -from mypy_django_plugin_newsemanal.lib import helpers +from mypy_django_plugin_newsemanal.lib import helpers, fullnames from mypy_django_plugin_newsemanal.transformers import fields from mypy_django_plugin_newsemanal.transformers.fields import get_field_descriptor_types @@ -35,8 +36,8 @@ class ModelClassInitializer(metaclass=ABCMeta): raise helpers.IncompleteDefnException(f'No {fullname!r} found') return sym.node - def lookup_field_typeinfo_or_incomplete_defn_error(self, field: Field) -> TypeInfo: - fullname = helpers.get_class_fullname(field.__class__) + def lookup_class_typeinfo_or_incomplete_defn_error(self, klass: type) -> TypeInfo: + fullname = helpers.get_class_fullname(klass) field_info = self.lookup_typeinfo_or_incomplete_defn_error(fullname) return field_info @@ -101,7 +102,7 @@ class AddRelatedModelsId(ModelClassInitializer): for field in model_cls._meta.get_fields(): if isinstance(field, ForeignKey): rel_primary_key_field = self.django_context.get_primary_key_field(field.related_model) - field_info = self.lookup_field_typeinfo_or_incomplete_defn_error(rel_primary_key_field) + field_info = self.lookup_class_typeinfo_or_incomplete_defn_error(rel_primary_key_field.__class__) is_nullable = self.django_context.fields_context.get_field_nullability(field, None) set_type, get_type = get_field_descriptor_types(field_info, is_nullable) self.add_new_node_to_model_class(field.attname, @@ -129,6 +130,22 @@ class AddManagers(ModelClassInitializer): default_manager = Instance(default_manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class('_default_manager', default_manager) + # add related managers + for relation in self.django_context.get_model_relations(model_cls): + attname = relation.related_name + if attname is None: + attname = relation.name + '_set' + + related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(relation.related_model) + if isinstance(relation, OneToOneRel): + self.add_new_node_to_model_class(attname, Instance(related_model_info, [])) + continue + if isinstance(relation, (ManyToOneRel, ManyToManyRel)): + manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS_FULLNAME) + self.add_new_node_to_model_class(attname, + Instance(manager_info, [Instance(related_model_info, [])])) + continue + def process_model_class(ctx: ClassDefContext, django_context: DjangoContext) -> None: diff --git a/test-data-newsemanal/typecheck/fields/test_related.yml b/test-data-newsemanal/typecheck/fields/test_related.yml index cc53fb3..4dcdfcc 100644 --- a/test-data-newsemanal/typecheck/fields/test_related.yml +++ b/test-data-newsemanal/typecheck/fields/test_related.yml @@ -18,13 +18,15 @@ publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE, related_name='books') -- case: test_foreign_key_field_creates_attribute_with_underscore_id +- case: foreign_key_field_creates_attribute_with_underscore_id main: | from myapp.models import Book book = Book() - reveal_type(book.publisher_id) # N: Revealed type is 'builtins.int' - reveal_type(book.owner_id) # N: Revealed type is 'Any' + reveal_type(book.publisher_id) # N: Revealed type is 'builtins.int*' + reveal_type(book.owner_id) # N: Revealed type is 'builtins.int*' installed_apps: + - django.contrib.contenttypes + - django.contrib.auth - myapp files: - path: myapp/__init__.py @@ -35,9 +37,9 @@ pass class Book(models.Model): publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE) - owner = models.ForeignKey(db_column='model_id', to='db.Unknown', on_delete=models.CASCADE) + owner = models.ForeignKey(db_column='model_id', to='auth.User', on_delete=models.CASCADE) -- case: test_foreign_key_field_different_order_of_params +- case: foreign_key_field_different_order_of_params main: | from myapp.models import Book, Publisher book = Book() @@ -61,7 +63,7 @@ related_name='books') publisher2 = models.ForeignKey(to=Publisher, related_name='books2', on_delete=models.CASCADE) -- case: test_to_parameter_as_string_with_application_name__model_imported +- case: to_parameter_as_string_with_application_name__model_imported main: | from myapp2.models import Book book = Book() @@ -85,6 +87,7 @@ - case: test_circular_dependency_in_imports_with_foreign_key main: | + from myapp import models installed_apps: - myapp files: @@ -113,16 +116,17 @@ reveal_type(View().app.views) # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp.models.View]' reveal_type(View().app.unknown) out: | - main:7: note: Revealed type is 'Any' - main:7: error: "App" has no attribute "unknown" + main:3: note: Revealed type is 'Any' + main:3: error: "App" has no attribute "unknown" installed_apps: - myapp + - myapp2 files: - path: myapp/__init__.py - path: myapp/models.py content: | from django.db import models - from myapp.models import App + from myapp2.models import App class View(models.Model): app = models.ForeignKey(to=App, related_name='views', on_delete=models.CASCADE) - path: myapp2/__init__.py @@ -150,12 +154,12 @@ class View(models.Model): app = models.ForeignKey(to=App, on_delete=models.CASCADE, related_name='views') class View2(View): - app = models.ForeignKey(to=App, on_delete=models.CASCADE, related_name='views2') + app2 = models.ForeignKey(to=App, on_delete=models.CASCADE, related_name='views2') - case: models_imported_inside_init_file_foreign_key main: | from myapp2.models import View - reveal_type(View().app.views) # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp.models.View]' + reveal_type(View().app.views) # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp2.models.View]' installed_apps: - myapp - myapp2 @@ -180,7 +184,7 @@ - case: models_imported_inside_init_file_one_to_one_field main: | from myapp2.models import Profile - reveal_type(Profile().user.profile) # N: Revealed type is 'myapp.models.Profile' + reveal_type(Profile().user.profile) # N: Revealed type is 'myapp2.models.Profile' installed_apps: - myapp - myapp2 @@ -189,7 +193,7 @@ - path: myapp/models/__init__.py content: | from .user import User - - path: myapp/models/app.py + - path: myapp/models/user.py content: | from django.db import models class User(models.Model): @@ -205,7 +209,7 @@ - case: models_triple_circular_reference main: | from myapp.models import App - reveal_type(App().owner) # N: Revealed type is 'myapp.models.user.User' + reveal_type(App().owner) # N: Revealed type is 'myapp.models.user.User*' reveal_type(App().owner.profile) # N: Revealed type is 'myapp.models.profile.Profile' installed_apps: - myapp @@ -224,7 +228,6 @@ - path: myapp/models/profile.py content: | from django.db import models - from myapp.models import User class Profile(models.Model): user = models.OneToOneField(to='myapp.User', related_name='profile', on_delete=models.CASCADE) - path: myapp/models/app.py @@ -253,7 +256,7 @@ - case: many_to_many_works_with_string_if_imported main: | from myapp.models import Member - reveal_type(Member().apps) # N: Revealed type is 'django.db.models.manager.RelatedManager*[myapp.models.App]' + reveal_type(Member().apps) # N: Revealed type is 'django.db.models.manager.RelatedManager*[myapp2.models.App]' installed_apps: - myapp - myapp2 @@ -262,9 +265,8 @@ - path: myapp/models.py content: | from django.db import models - from myapp.models import App class Member(models.Model): - apps = models.ManyToManyField(to='myapp.App', related_name='members') + apps = models.ManyToManyField(to='myapp2.App', related_name='members') - path: myapp2/__init__.py - path: myapp2/models.py content: | @@ -335,15 +337,17 @@ - case: underscore_id_attribute_has_set_type_of_primary_key_if_explicit main: | - from myapp.models import Book - reveal_type(Book().publisher_id) # N: Revealed type is 'builtins.str' + import datetime + from myapp.models import Book, Book2 + + reveal_type(Book().publisher_id) # N: Revealed type is 'builtins.str*' Book(publisher_id=1) Book(publisher_id='hello') Book(publisher_id=datetime.datetime.now()) # E: Incompatible type for "publisher_id" of "Book" (got "datetime", expected "Union[str, int, Combinable, None]") Book.objects.create(publisher_id=1) Book.objects.create(publisher_id='hello') - reveal_type(Book2().publisher_id) # N: Revealed type is 'builtins.int' + reveal_type(Book2().publisher_id) # N: Revealed type is 'builtins.int*' Book2(publisher_id=1) Book2(publisher_id=[]) # E: Incompatible type for "publisher_id" of "Book2" (got "List[Any]", expected "Union[float, int, str, Combinable, None]") Book2.objects.create(publisher_id=1) @@ -392,8 +396,8 @@ reveal_type(publisher.books) reveal_type(publisher.books2) # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp.models.Book]' out: | - main:16: note: Revealed type is 'Any' - main:16: error: "Publisher" has no attribute "books"; maybe "books2"? + main:6: error: "Publisher" has no attribute "books"; maybe "books2"? + main:6: note: Revealed type is 'Any' installed_apps: - myapp files: diff --git a/test-data-newsemanal/typecheck/models/test_inheritance.yml b/test-data-newsemanal/typecheck/models/test_inheritance.yml index a933dc4..d31f16b 100644 --- a/test-data-newsemanal/typecheck/models/test_inheritance.yml +++ b/test-data-newsemanal/typecheck/models/test_inheritance.yml @@ -55,9 +55,9 @@ reveal_type(b_instance.non_existent_attribute) b_instance.non_existent_attribute = 2 out: | - main:10: note: Revealed type is 'Any' - main:10: error: "B" has no attribute "non_existent_attribute" - main:11: error: "B" has no attribute "non_existent_attribute" + main:5: note: Revealed type is 'Any' + main:5: error: "B" has no attribute "non_existent_attribute" + main:6: error: "B" has no attribute "non_existent_attribute" installed_apps: - myapp files: