Fix ForeignKey type for self-reference defined in the abstract model (#200)

This commit is contained in:
Maxim Kurnikov
2019-10-05 21:36:29 +03:00
committed by GitHub
parent db9ff6aaf6
commit 7e3f4bfa02
6 changed files with 82 additions and 16 deletions

View File

@@ -37,6 +37,14 @@ def _get_current_field_from_assignment(ctx: FunctionContext, django_context: Dja
return current_field
def reparametrize_related_field_type(related_field_type: Instance, set_type, get_type) -> Instance:
args = [
helpers.convert_any_to_type(related_field_type.args[0], set_type),
helpers.convert_any_to_type(related_field_type.args[1], get_type),
]
return helpers.reparametrize_instance(related_field_type, new_args=args)
def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
current_field = _get_current_field_from_assignment(ctx, django_context)
if current_field is None:
@@ -48,6 +56,25 @@ def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context
if related_model_cls is None:
return AnyType(TypeOfAny.from_error)
default_related_field_type = set_descriptor_types_for_field(ctx)
# self reference with abstract=True on the model where ForeignKey is defined
current_model_cls = current_field.model
if (current_model_cls._meta.abstract
and current_model_cls == related_model_cls):
# for all derived non-abstract classes, set variable with this name to
# __get__/__set__ of ForeignKey of derived model
for model_cls in django_context.all_registered_model_classes:
if issubclass(model_cls, current_model_cls) and not model_cls._meta.abstract:
derived_model_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), model_cls)
if derived_model_info is not None:
fk_ref_type = Instance(derived_model_info, [])
derived_fk_type = reparametrize_related_field_type(default_related_field_type,
set_type=fk_ref_type, get_type=fk_ref_type)
helpers.add_new_sym_for_info(derived_model_info,
name=current_field.name,
sym_type=derived_fk_type)
related_model = related_model_cls
related_model_to_set = related_model_cls
if related_model_to_set._meta.proxy_for_model is not None:
@@ -69,13 +96,10 @@ def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context
else:
related_model_to_set_type = Instance(related_model_to_set_info, []) # type: ignore
default_related_field_type = set_descriptor_types_for_field(ctx)
# replace Any with referred_to_type
args = [
helpers.convert_any_to_type(default_related_field_type.args[0], related_model_to_set_type),
helpers.convert_any_to_type(default_related_field_type.args[1], related_model_type),
]
return helpers.reparametrize_instance(default_related_field_type, new_args=args)
return reparametrize_related_field_type(default_related_field_type,
set_type=related_model_to_set_type,
get_type=related_model_type)
def get_field_descriptor_types(field_info: TypeInfo, is_nullable: bool) -> Tuple[MypyType, MypyType]:

View File

@@ -7,9 +7,7 @@ from django.db.models.fields.related import ForeignKey
from django.db.models.fields.reverse_related import (
ManyToManyRel, ManyToOneRel, OneToOneRel,
)
from mypy.nodes import (
ARG_STAR2, MDEF, Argument, Context, SymbolTableNode, TypeInfo, Var,
)
from mypy.nodes import ARG_STAR2, Argument, Context, TypeInfo, Var
from mypy.plugin import ClassDefContext
from mypy.plugins import common
from mypy.types import AnyType, Instance
@@ -51,8 +49,9 @@ class ModelClassInitializer:
return var
def add_new_node_to_model_class(self, name: str, typ: MypyType) -> None:
var = self.create_new_var(name, typ)
self.model_classdef.info.names[name] = SymbolTableNode(MDEF, var, plugin_generated=True)
helpers.add_new_sym_for_info(self.model_classdef.info,
name=name,
sym_type=typ)
def run(self) -> None:
model_cls = self.django_context.get_model_class_by_fullname(self.model_classdef.fullname)
@@ -114,6 +113,9 @@ class AddRelatedModelsId(ModelClassInitializer):
AnyType(TypeOfAny.explicit))
continue
if related_model_cls._meta.abstract:
continue
rel_primary_key_field = self.django_context.get_primary_key_field(related_model_cls)
field_info = self.lookup_class_typeinfo_or_incomplete_defn_error(rel_primary_key_field.__class__)
is_nullable = self.django_context.get_field_nullability(field, None)