mirror of
https://github.com/davidhalter/django-stubs.git
synced 2025-12-07 20:54:29 +08:00
Fix ForeignKey type for self-reference defined in the abstract model (#200)
This commit is contained in:
@@ -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]:
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user