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_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_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 set_auth_user_model_as_type_for_request_user(ctx: AttributeContext, django_context: DjangoContext) -> MypyType:
# Imported here because django isn't properly loaded yet when module is loaded
from django.contrib.auth.base_user import AbstractBaseUser
from django.contrib.auth.models import AnonymousUser
abstract_base_user_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), AbstractBaseUser)
anonymous_user_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), AnonymousUser)
# This shouldn't be able to happen, as we managed to import the models above.
assert abstract_base_user_info is not None
assert anonymous_user_info is not None
if ctx.default_attr_type != UnionType([Instance(abstract_base_user_info, []), Instance(anonymous_user_info, [])]):
# Type has been changed from the default in django-stubs.
# I.e. HttpRequest has been subclassed and user-type overridden, so let's leave it as is.
return ctx.default_attr_type
auth_user_model = django_context.settings.AUTH_USER_MODEL
user_cls = django_context.apps_registry.get_model(auth_user_model)
user_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), user_cls)
if user_info is None:
return ctx.default_attr_type
return UnionType([Instance(user_info, []), Instance(anonymous_user_info, [])])
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.get_field_related_model_cls(current_field)
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:
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
# replace Any with referred_to_type
return reparametrize_related_field_type(default_related_field_type,
set_type=related_model_to_set_type,
get_type=related_model_type)
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 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 set_auth_user_model_as_type_for_request_user(
ctx: AttributeContext, django_context: DjangoContext) -> MypyType:
auth_user_model = django_context.settings.AUTH_USER_MODEL
model_cls = django_context.apps_registry.get_model(auth_user_model)
model_info = helpers.lookup_class_typeinfo(
helpers.get_typechecker_api(ctx), model_cls)
if model_info is None:
return ctx.default_attr_type
return Instance(model_info, [])
def set_auth_user_model_as_type_for_request_user(
ctx: AttributeContext, django_context: DjangoContext) -> MypyType:
auth_user_model = django_context.settings.AUTH_USER_MODEL
user_cls = django_context.apps_registry.get_model(auth_user_model)
user_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx),
user_cls)
if user_info is None:
return ctx.default_attr_type
# Imported here because django isn't properly loaded yet when module is loaded
from django.contrib.auth.models import AnonymousUser
anonymous_user_info = helpers.lookup_class_typeinfo(
helpers.get_typechecker_api(ctx), AnonymousUser)
if anonymous_user_info is None:
# This shouldn't be able to happen, as we managed to import the model above...
return Instance(user_info, [])
return UnionType(
[Instance(user_info, []),
Instance(anonymous_user_info, [])])
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 set_auth_user_model_boolean_fields(
ctx: AttributeContext, django_context: DjangoContext) -> MypyType:
boolinfo = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx),
bool)
assert boolinfo is not None
return Instance(boolinfo, [])
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
