add related managers support

This commit is contained in:
Maxim Kurnikov
2019-07-16 18:05:23 +03:00
parent 2cb1f257eb
commit 7b1b1b6bfe
6 changed files with 69 additions and 40 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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