VimPlug/django-stubs
1
0
Fork 0
mirror of https://github.com/davidhalter/django-stubs.git synced 2025-12-16 00:37:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

add related managers support

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

6
mypy_django_plugin_newsemanal/django/context.py
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):

14
mypy_django_plugin_newsemanal/main.py
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)

8
mypy_django_plugin_newsemanal/transformers/fields.py
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)

25
mypy_django_plugin_newsemanal/transformers/models.py
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: