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 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 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:
return AnyType(TypeOfAny.from_error)
assert isinstance(current_field, RelatedField)
related_model_cls = django_context.fields_context.get_related_model_cls(
current_field)
related_model = related_model_cls
related_model_to_set = related_model_cls
if related_model_to_set._meta.proxy_for_model is not None:
related_model_to_set = related_model_to_set._meta.proxy_for_model
typechecker_api = helpers.get_typechecker_api(ctx)
related_model_info = helpers.lookup_class_typeinfo(typechecker_api,
related_model)
if related_model_info is None:
# maybe no type stub
related_model_type = AnyType(TypeOfAny.unannotated)
else:
related_model_type = Instance(related_model_info, []) # type: ignore
related_model_to_set_info = helpers.lookup_class_typeinfo(
typechecker_api, related_model_to_set)
if related_model_to_set_info is None:
# maybe no type stub
related_model_to_set_type = AnyType(TypeOfAny.unannotated)
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)
def determine_type_of_array_field(ctx: FunctionContext,
django_context: DjangoContext) -> MypyType:
default_return_type = set_descriptor_types_for_field(ctx)
base_field_arg_type = helpers.get_call_argument_type_by_name(
ctx, "base_field")
if not base_field_arg_type or not isinstance(base_field_arg_type,
Instance):
return default_return_type
def drop_combinable(_type: MypyType) -> Optional[MypyType]:
if isinstance(_type, Instance) and _type.type.has_base(
fullnames.COMBINABLE_EXPRESSION_FULLNAME):
return None
elif isinstance(_type, UnionType):
items_without_combinable = []
for item in _type.items:
reduced = drop_combinable(item)
if reduced is not None:
items_without_combinable.append(reduced)
if len(items_without_combinable) > 1:
return UnionType(
items_without_combinable,
line=_type.line,
column=_type.column,
is_evaluated=_type.is_evaluated,
uses_pep604_syntax=_type.uses_pep604_syntax,
)
elif len(items_without_combinable) == 1:
return items_without_combinable[0]
else:
return None
return _type
# Both base_field and return type should derive from Field and thus expect 2 arguments
assert len(base_field_arg_type.args) == len(default_return_type.args) == 2
args = []
for new_type, default_arg in zip(base_field_arg_type.args,
default_return_type.args):
# Drop any base_field Combinable type
reduced = drop_combinable(new_type)
if reduced is None:
ctx.api.fail(
f"Can't have ArrayField expecting {fullnames.COMBINABLE_EXPRESSION_FULLNAME!r} as data type",
ctx.context,
)
else:
new_type = reduced
args.append(helpers.convert_any_to_type(default_arg, new_type))
return helpers.reparametrize_instance(default_return_type, args)
def determine_type_of_array_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
default_return_type = set_descriptor_types_for_field(ctx)
base_field_arg_type = helpers.get_call_argument_type_by_name(ctx, 'base_field')
if not base_field_arg_type or not isinstance(base_field_arg_type, Instance):
return default_return_type
base_type = base_field_arg_type.args[1] # extract __get__ type
args = []
for default_arg in default_return_type.args:
args.append(helpers.convert_any_to_type(default_arg, base_type))
return helpers.reparametrize_instance(default_return_type, args)
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
