def get_expected_types(self, api: TypeChecker, model_cls: Type[Model], *, method: str) -> Dict[str, MypyType]:
contenttypes_in_apps = self.apps_registry.is_installed("django.contrib.contenttypes")
if contenttypes_in_apps:
from django.contrib.contenttypes.fields import GenericForeignKey
expected_types = {}
# add pk if not abstract=True
if not model_cls._meta.abstract:
primary_key_field = self.get_primary_key_field(model_cls)
field_set_type = self.get_field_set_type(api, primary_key_field, method=method)
expected_types["pk"] = field_set_type
for field in model_cls._meta.get_fields():
if isinstance(field, Field):
field_name = field.attname
field_set_type = self.get_field_set_type(api, field, method=method)
expected_types[field_name] = field_set_type
if isinstance(field, ForeignKey):
field_name = field.name
foreign_key_info = helpers.lookup_class_typeinfo(api, field.__class__)
if foreign_key_info is None:
# maybe there's no type annotation for the field
expected_types[field_name] = AnyType(TypeOfAny.unannotated)
continue
related_model = self.get_field_related_model_cls(field)
if related_model is None:
expected_types[field_name] = AnyType(TypeOfAny.from_error)
continue
if related_model._meta.proxy_for_model is not None:
related_model = related_model._meta.proxy_for_model
related_model_info = helpers.lookup_class_typeinfo(api, related_model)
if related_model_info is None:
expected_types[field_name] = AnyType(TypeOfAny.unannotated)
continue
is_nullable = self.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 contenttypes_in_apps and isinstance(field, GenericForeignKey):
# it's generic, so cannot set specific model
field_name = field.name
gfk_info = helpers.lookup_class_typeinfo(api, field.__class__)
if gfk_info is None:
gfk_set_type = AnyType(TypeOfAny.unannotated)
else:
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
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 = {}
# 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=method)
expected_types['pk'] = 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=method)
expected_types[field_name] = field_set_type
if isinstance(field, ForeignKey):
field_name = field.name
foreign_key_info = helpers.lookup_class_typeinfo(
api, field.__class__)
if foreign_key_info is None:
# maybe there's no type annotation for the field
expected_types[field_name] = AnyType(
TypeOfAny.unannotated)
continue
related_model = field.related_model
if related_model._meta.proxy_for_model:
related_model = field.related_model._meta.proxy_for_model
related_model_info = helpers.lookup_class_typeinfo(
api, related_model)
if related_model_info is None:
expected_types[field_name] = AnyType(
TypeOfAny.unannotated)
continue
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
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
def get_field_descriptor_types(field_info: TypeInfo,
is_nullable: bool) -> Tuple[MypyType, MypyType]:
set_type = helpers.get_private_descriptor_type(field_info,
'_pyi_private_set_type',
is_nullable=is_nullable)
get_type = helpers.get_private_descriptor_type(field_info,
'_pyi_private_get_type',
is_nullable=is_nullable)
return set_type, get_type
def get_field_get_type(self, api: TypeChecker, field: Field, *,
method: str) -> MypyType:
""" Get a type of __get__ for this specific Django field. """
field_info = helpers.lookup_class_typeinfo(api, field.__class__)
if field_info is None:
return AnyType(TypeOfAny.unannotated)
is_nullable = self.get_field_nullability(field, method)
if isinstance(field, RelatedField):
related_model_cls = self.get_field_related_model_cls(field)
if related_model_cls is None:
return AnyType(TypeOfAny.from_error)
if method == 'values':
primary_key_field = self.get_primary_key_field(
related_model_cls)
return self.get_field_get_type(api,
primary_key_field,
method=method)
model_info = helpers.lookup_class_typeinfo(api, related_model_cls)
if model_info is None:
return AnyType(TypeOfAny.unannotated)
return Instance(model_info, [])
else:
return helpers.get_private_descriptor_type(field_info,
'_pyi_private_get_type',
is_nullable=is_nullable)
def get_field_lookup_exact_type(self, api: TypeChecker,
field: Field) -> MypyType:
if isinstance(field, (RelatedField, ForeignObjectRel)):
related_model_cls = field.related_model
primary_key_field = self.get_primary_key_field(related_model_cls)
primary_key_type = self.get_field_get_type(api,
primary_key_field,
method='init')
rel_model_info = helpers.lookup_class_typeinfo(
api, related_model_cls)
if rel_model_info is None:
return AnyType(TypeOfAny.explicit)
model_and_primary_key_type = UnionType.make_union(
[Instance(rel_model_info, []), primary_key_type])
return helpers.make_optional(model_and_primary_key_type)
# return helpers.make_optional(Instance(rel_model_info, []))
field_info = helpers.lookup_class_typeinfo(api, field.__class__)
if field_info is None:
return AnyType(TypeOfAny.explicit)
return helpers.get_private_descriptor_type(field_info,
'_pyi_lookup_exact_type',
is_nullable=field.null)
def resolve_lookup_expected_type(self, ctx: MethodContext,
model_cls: Type[Model],
lookup: str) -> MypyType:
query = Query(model_cls)
try:
lookup_parts, field_parts, is_expression = query.solve_lookup_type(
lookup)
if is_expression:
return AnyType(TypeOfAny.explicit)
except FieldError as exc:
ctx.api.fail(exc.args[0], ctx.context)
return AnyType(TypeOfAny.from_error)
field = self._resolve_field_from_parts(field_parts, model_cls)
lookup_cls = None
if lookup_parts:
lookup = lookup_parts[-1]
lookup_cls = field.get_lookup(lookup)
if lookup_cls is None:
# unknown lookup
return AnyType(TypeOfAny.explicit)
if lookup_cls is None or isinstance(lookup_cls, Exact):
return self.get_field_lookup_exact_type(
helpers.get_typechecker_api(ctx), field)
assert lookup_cls is not None
lookup_info = helpers.lookup_class_typeinfo(
helpers.get_typechecker_api(ctx), lookup_cls)
if lookup_info is None:
return AnyType(TypeOfAny.explicit)
for lookup_base in helpers.iter_bases(lookup_info):
if lookup_base.args and isinstance(lookup_base.args[0], Instance):
lookup_type: MypyType = lookup_base.args[0]
# if it's Field, consider lookup_type a __get__ of current field
if (isinstance(lookup_type, Instance)
and lookup_type.type.fullname()
== fullnames.FIELD_FULLNAME):
field_info = helpers.lookup_class_typeinfo(
helpers.get_typechecker_api(ctx), field.__class__)
if field_info is None:
return AnyType(TypeOfAny.explicit)
lookup_type = helpers.get_private_descriptor_type(
field_info,
'_pyi_private_get_type',
is_nullable=field.null)
return lookup_type
return AnyType(TypeOfAny.explicit)
def get_field_set_type(self, api: TypeChecker, field: Field, *, method: str) -> MypyType:
""" Get a type of __set__ for this specific Django field. """
target_field = field
if isinstance(field, ForeignKey):
target_field = field.target_field
field_info = helpers.lookup_class_typeinfo(api, target_field.__class__)
if field_info is None:
return AnyType(TypeOfAny.from_error)
field_set_type = helpers.get_private_descriptor_type(field_info, '_pyi_private_set_type',
is_nullable=self.get_field_nullability(field, method))
if isinstance(target_field, ArrayField):
argument_field_type = self.get_field_set_type(api, target_field.base_field, method=method)
field_set_type = helpers.convert_any_to_type(field_set_type, argument_field_type)
return field_set_type
def get_expected_types(self, api: TypeChecker, model_cls: Type[Model], *,
method: str) -> Dict[str, MypyType]:
contenttypes_in_apps = self.apps_registry.is_installed(
"django.contrib.contenttypes")
if contenttypes_in_apps:
from django.contrib.contenttypes.fields import GenericForeignKey
expected_types = {}
# add pk if not abstract=True
if not model_cls._meta.abstract:
primary_key_field = self.get_primary_key_field(model_cls)
field_set_type = self.get_field_set_type(api,
primary_key_field,
method=method)
expected_types["pk"] = field_set_type
def get_field_set_type_from_model_type_info(
info: Optional[TypeInfo],
field_name: str) -> Optional[MypyType]:
if info is None:
return None
field_node = info.names.get(field_name)
if field_node is None or not isinstance(field_node.type, Instance):
return None
elif not field_node.type.args:
# Field declares a set and a get type arg. Fallback to `None` when we can't find any args
return None
set_type = field_node.type.args[0]
return set_type
model_info = helpers.lookup_class_typeinfo(api, model_cls)
for field in model_cls._meta.get_fields():
if isinstance(field, Field):
field_name = field.attname
# Try to retrieve set type from a model's TypeInfo object and fallback to retrieving it manually
# from django-stubs own declaration. This is to align with the setter types declared for
# assignment.
field_set_type = get_field_set_type_from_model_type_info(
model_info, field_name) or self.get_field_set_type(
api, field, method=method)
expected_types[field_name] = field_set_type
if isinstance(field, ForeignKey):
field_name = field.name
foreign_key_info = helpers.lookup_class_typeinfo(
api, field.__class__)
if foreign_key_info is None:
# maybe there's no type annotation for the field
expected_types[field_name] = AnyType(
TypeOfAny.unannotated)
continue
related_model = self.get_field_related_model_cls(field)
if related_model is None:
expected_types[field_name] = AnyType(
TypeOfAny.from_error)
continue
if related_model._meta.proxy_for_model is not None:
related_model = related_model._meta.proxy_for_model
related_model_info = helpers.lookup_class_typeinfo(
api, related_model)
if related_model_info is None:
expected_types[field_name] = AnyType(
TypeOfAny.unannotated)
continue
is_nullable = self.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 contenttypes_in_apps and isinstance(field, GenericForeignKey):
# it's generic, so cannot set specific model
field_name = field.name
gfk_info = helpers.lookup_class_typeinfo(api, field.__class__)
if gfk_info is None:
gfk_set_type: MypyType = AnyType(TypeOfAny.unannotated)
else:
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
