add related managers support

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):

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:
+            # forward relations
-                related_model_module = related_object.related_model.__module__
+            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)

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)
 

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:
+    def lookup_class_typeinfo_or_incomplete_defn_error(self, klass: type) -> TypeInfo:
-        fullname = helpers.get_class_fullname(field.__class__)
+        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:

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.publisher_id)  # N: Revealed type is 'builtins.int*'
-        reveal_type(book.owner_id)  # N: Revealed type is 'Any'
+        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:3: note: Revealed type is 'Any'
-        main:7: error: "App" has no attribute "unknown"
+        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
+        import datetime
-        reveal_type(Book().publisher_id)  # N: Revealed type is 'builtins.str'
+        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:6: error: "Publisher" has no attribute "books"; maybe "books2"?
-        main:16: error: "Publisher" has no attribute "books"; maybe "books2"?
+        main:6: note: Revealed type is 'Any'
     installed_apps:
         - myapp
     files:

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:5: note: Revealed type is 'Any'
-        main:10: error: "B" has no attribute "non_existent_attribute"
+        main:5: error: "B" has no attribute "non_existent_attribute"
-        main:11: error: "B" has no attribute "non_existent_attribute"
+        main:6: error: "B" has no attribute "non_existent_attribute"
     installed_apps:
         - myapp
     files: