add GenericForeignKey support, remove some false-positives

This commit is contained in:
Maxim Kurnikov
2019-07-18 18:31:37 +03:00
parent bfa77efef5
commit f2e79d3bfb
15 changed files with 111 additions and 62 deletions

View File

@@ -1,11 +1,13 @@
import os
from collections import defaultdict
from dataclasses import dataclass
from typing import Any, Dict, Iterator, List, Optional, TYPE_CHECKING, Tuple, Type, Sequence
from typing import Any, Dict, Iterator, List, Optional, TYPE_CHECKING, Tuple, Type
from django.core.exceptions import FieldError, FieldDoesNotExist
from django.core.exceptions import FieldError
from django.db.models.base import Model
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 Instance, Type as MypyType
@@ -13,9 +15,6 @@ 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, ManyToOneRel, ManyToManyRel
from django.db.models.sql.query import Query
from mypy_django_plugin.lib import helpers
if TYPE_CHECKING:
@@ -99,7 +98,7 @@ class DjangoFieldsContext:
return Instance(model_info, [])
else:
return helpers.get_private_descriptor_type(field_info, '_pyi_private_get_type',
is_nullable=is_nullable)
is_nullable=is_nullable)
class DjangoLookupsContext:
@@ -184,27 +183,38 @@ class DjangoContext:
raise ValueError('No primary key defined')
def get_expected_types(self, api: TypeChecker, model_cls: Type[Model], method: str) -> Dict[str, MypyType]:
from django.contrib.contenttypes.fields import GenericForeignKey
expected_types = {}
if method == '__init__':
# add pk
primary_key_field = self.get_primary_key_field(model_cls)
field_set_type = self.fields_context.get_field_set_type(api, primary_key_field, method)
expected_types['pk'] = field_set_type
# add pk
primary_key_field = self.get_primary_key_field(model_cls)
field_set_type = self.fields_context.get_field_set_type(api, primary_key_field, method)
expected_types['pk'] = field_set_type
for field in self.get_model_fields(model_cls):
field_name = field.attname
field_set_type = self.fields_context.get_field_set_type(api, field, method)
expected_types[field_name] = field_set_type
for field in model_cls._meta.get_fields():
if isinstance(field, Field):
field_name = field.attname
field_set_type = self.fields_context.get_field_set_type(api, field, method)
expected_types[field_name] = field_set_type
if isinstance(field, ForeignKey):
if isinstance(field, ForeignKey):
field_name = field.name
foreign_key_info = helpers.lookup_class_typeinfo(api, field.__class__)
related_model_info = helpers.lookup_class_typeinfo(api, field.related_model)
is_nullable = self.fields_context.get_field_nullability(field, method)
foreign_key_set_type = helpers.get_private_descriptor_type(foreign_key_info,
'_pyi_private_set_type',
is_nullable=is_nullable)
model_set_type = helpers.convert_any_to_type(foreign_key_set_type,
Instance(related_model_info, []))
expected_types[field_name] = model_set_type
elif isinstance(field, GenericForeignKey):
# it's generic, so cannot set specific model
field_name = field.name
foreign_key_info = helpers.lookup_class_typeinfo(api, field.__class__)
related_model_info = helpers.lookup_class_typeinfo(api, field.related_model)
is_nullable = self.fields_context.get_field_nullability(field, method)
foreign_key_set_type = helpers.get_private_descriptor_type(foreign_key_info,
'_pyi_private_set_type',
is_nullable=is_nullable)
model_set_type = helpers.convert_any_to_type(foreign_key_set_type,
Instance(related_model_info, []))
expected_types[field_name] = model_set_type
gfk_info = helpers.lookup_class_typeinfo(api, field.__class__)
gfk_set_type = helpers.get_private_descriptor_type(gfk_info, '_pyi_private_set_type',
is_nullable=True)
expected_types[field_name] = gfk_set_type
return expected_types

View File

@@ -1,6 +1,6 @@
from typing import Optional, Tuple, cast
from mypy.nodes import TypeInfo
from mypy.nodes import MypyFile, TypeInfo
from mypy.plugin import FunctionContext
from mypy.types import AnyType, CallableType, Instance, Type as MypyType, TypeOfAny
@@ -14,11 +14,13 @@ def get_referred_to_model_fullname(ctx: FunctionContext, django_context: DjangoC
assert isinstance(to_arg_type.ret_type, Instance)
return to_arg_type.ret_type.type.fullname()
outer_model_info = ctx.api.tscope.classes[-1]
outer_model_info = ctx.api.scope.active_class()
if not outer_model_info or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME):
# not inside models.Model class
return None
assert isinstance(outer_model_info, TypeInfo)
to_arg_expr = helpers.get_call_argument_by_name(ctx, 'to')
model_string = helpers.resolve_string_attribute_value(to_arg_expr, ctx, django_context)
if model_string is None:
# unresolvable
@@ -28,10 +30,21 @@ def get_referred_to_model_fullname(ctx: FunctionContext, django_context: DjangoC
return outer_model_info.fullname()
if '.' not in model_string:
# same file class
current_module = ctx.api.tree
if model_string not in current_module.names:
model_cls_is_accessible = False
for scope in ctx.api.scope.stack:
if isinstance(scope, (MypyFile, TypeInfo)):
model_class_candidate = scope.names.get(model_string)
model_cls_is_accessible = (model_class_candidate is not None
and isinstance(model_class_candidate.node, TypeInfo)
and model_class_candidate.node.has_base(fullnames.MODEL_CLASS_FULLNAME))
if model_cls_is_accessible:
break
# TODO: FuncItem
if not model_cls_is_accessible:
ctx.api.fail(f'No model {model_string!r} defined in the current module', ctx.context)
return None
return outer_model_info.module_name + '.' + model_string
app_label, model_name = model_string.split('.')
@@ -88,12 +101,6 @@ def transform_into_proper_return_type(ctx: FunctionContext, django_context: Djan
default_return_type = ctx.default_return_type
assert isinstance(default_return_type, Instance)
# bail out if we're inside migration, not supported yet
active_class = ctx.api.scope.active_class()
if active_class is not None:
if active_class.has_base(fullnames.MIGRATION_CLASS_FULLNAME):
return ctx.default_return_type
if helpers.has_any_of_bases(default_return_type.type, fullnames.RELATED_FIELDS_CLASSES):
return fill_descriptor_types_for_related_field(ctx, django_context)