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 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 _get_current_field_from_assignment(
ctx: FunctionContext,
django_context: DjangoContext) -> Optional[Field]:
outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
if (outer_model_info is None or not helpers.is_model_subclass_info(
outer_model_info, django_context)):
return None
field_name = None
for stmt in outer_model_info.defn.defs.body:
if isinstance(stmt, AssignmentStmt):
if stmt.rvalue == ctx.context:
if not isinstance(stmt.lvalues[0], NameExpr):
return None
field_name = stmt.lvalues[0].name
break
if field_name is None:
return None
model_cls = django_context.get_model_class_by_fullname(
outer_model_info.fullname)
if model_cls is None:
return None
current_field = model_cls._meta.get_field(field_name)
return current_field
def get_field_type_from_lookup(
ctx: MethodContext,
django_context: DjangoContext,
model_cls: Type[Model],
*,
method: str,
lookup: str,
silent_on_error: bool = False,
) -> Optional[MypyType]:
try:
lookup_field = django_context.resolve_lookup_into_field(
model_cls, lookup)
except FieldError as exc:
if not silent_on_error:
ctx.api.fail(exc.args[0], ctx.context)
return None
except LookupsAreUnsupported:
return AnyType(TypeOfAny.explicit)
if (isinstance(lookup_field, RelatedField)
and lookup_field.column == lookup) or isinstance(
lookup_field, ForeignObjectRel):
related_model_cls = django_context.get_field_related_model_cls(
lookup_field)
if related_model_cls is None:
return AnyType(TypeOfAny.from_error)
lookup_field = django_context.get_primary_key_field(related_model_cls)
field_get_type = django_context.get_field_get_type(
helpers.get_typechecker_api(ctx), lookup_field, method=method)
return field_get_type
def get_type_of_settings_attribute(ctx: AttributeContext,
django_context: DjangoContext) -> MypyType:
assert isinstance(ctx.context, MemberExpr)
setting_name = ctx.context.name
if not hasattr(django_context.settings, setting_name):
ctx.api.fail(f"'Settings' object has no attribute {setting_name!r}",
ctx.context)
return ctx.default_attr_type
typechecker_api = helpers.get_typechecker_api(ctx)
# first look for the setting in the project settings file, then global settings
settings_module = typechecker_api.modules.get(
django_context.django_settings_module)
global_settings_module = typechecker_api.modules.get(
'django.conf.global_settings')
for module in [settings_module, global_settings_module]:
if module is not None:
sym = module.names.get(setting_name)
if sym is not None and sym.type is not None:
return sym.type
# if by any reason it isn't present there, get type from django settings
value = getattr(django_context.settings, setting_name)
value_fullname = helpers.get_class_fullname(value.__class__)
value_info = helpers.lookup_fully_qualified_typeinfo(
typechecker_api, value_fullname)
if value_info is None:
return ctx.default_attr_type
return Instance(value_info, [])
def _get_field_instance(ctx: MethodContext, field_fullname: str) -> MypyType:
field_info = helpers.lookup_fully_qualified_typeinfo(
helpers.get_typechecker_api(ctx), field_fullname)
if field_info is None:
return AnyType(TypeOfAny.unannotated)
return Instance(field_info,
[AnyType(TypeOfAny.explicit),
AnyType(TypeOfAny.explicit)])
def resolve_manager_method_from_instance(instance: Instance, method_name: str,
ctx: AttributeContext) -> MypyType:
api = helpers.get_typechecker_api(ctx)
method_type = get_method_type_from_dynamic_manager(
api, method_name,
instance.type) or get_method_type_from_reverse_manager(
api, method_name, instance.type)
return method_type if method_type is not None else ctx.default_attr_type
def determine_proper_manager_type(ctx: FunctionContext) -> MypyType:
default_return_type = ctx.default_return_type
assert isinstance(default_return_type, Instance)
outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
if (outer_model_info is None
or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME)):
return default_return_type
return helpers.reparametrize_instance(default_return_type, [Instance(outer_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
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 get_user_model_hook(ctx: FunctionContext,
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_cls_fullname = helpers.get_class_fullname(model_cls)
model_info = helpers.lookup_fully_qualified_typeinfo(
helpers.get_typechecker_api(ctx), model_cls_fullname)
if model_info is None:
return AnyType(TypeOfAny.unannotated)
return TypeType(Instance(model_info, []))
def transform_into_proper_return_type(
ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
default_return_type = ctx.default_return_type
assert isinstance(default_return_type, Instance)
outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
if (outer_model_info is None or
not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME) and
outer_model_info.fullname() not in helpers.get_all_model_mixins(
helpers.get_typechecker_api(ctx))):
# not inside models.Model class
return ctx.default_return_type
assert isinstance(outer_model_info, TypeInfo)
if helpers.has_any_of_bases(default_return_type.type,
fullnames.RELATED_FIELDS_CLASSES):
return fill_descriptor_types_for_related_field(ctx, django_context)
if default_return_type.type.has_base(fullnames.ARRAY_FIELD_FULLNAME):
return determine_type_of_array_field(ctx, django_context)
return set_descriptor_types_for_field(ctx)
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_type_from_lookup(ctx: MethodContext, django_context: DjangoContext, model_cls: Type[Model],
*, method: str, lookup: str) -> Optional[MypyType]:
try:
lookup_field = django_context.lookups_context.resolve_lookup(model_cls, lookup)
except FieldError as exc:
ctx.api.fail(exc.args[0], ctx.context)
return None
if isinstance(lookup_field, RelatedField) and lookup_field.column == lookup:
lookup_field = django_context.get_primary_key_field(lookup_field.related_model)
field_get_type = django_context.fields_context.get_field_get_type(helpers.get_typechecker_api(ctx),
lookup_field, method=method)
return field_get_type
def get_values_list_row_type(ctx: MethodContext, django_context: DjangoContext, model_cls: Type[Model],
flat: bool, named: bool) -> MypyType:
field_lookups = resolve_field_lookups(ctx.args[0], django_context)
if field_lookups is None:
return AnyType(TypeOfAny.from_error)
typechecker_api = helpers.get_typechecker_api(ctx)
if len(field_lookups) == 0:
if flat:
primary_key_field = django_context.get_primary_key_field(model_cls)
lookup_type = get_field_type_from_lookup(ctx, django_context, model_cls,
lookup=primary_key_field.attname, method='values_list')
assert lookup_type is not None
return lookup_type
elif named:
column_types: 'OrderedDict[str, MypyType]' = OrderedDict()
for field in django_context.get_model_fields(model_cls):
column_type = django_context.get_field_get_type(typechecker_api, field,
method='values_list')
column_types[field.attname] = column_type
return helpers.make_oneoff_named_tuple(typechecker_api, 'Row', column_types)
else:
# flat=False, named=False, all fields
field_lookups = []
for field in django_context.get_model_fields(model_cls):
field_lookups.append(field.attname)
if len(field_lookups) > 1 and flat:
typechecker_api.fail("'flat' is not valid when 'values_list' is called with more than one field", ctx.context)
return AnyType(TypeOfAny.from_error)
column_types = OrderedDict()
for field_lookup in field_lookups:
lookup_field_type = get_field_type_from_lookup(ctx, django_context, model_cls,
lookup=field_lookup, method='values_list')
if lookup_field_type is None:
return AnyType(TypeOfAny.from_error)
column_types[field_lookup] = lookup_field_type
if flat:
assert len(column_types) == 1
row_type = next(iter(column_types.values()))
elif named:
row_type = helpers.make_oneoff_named_tuple(typechecker_api, 'Row', column_types)
else:
row_type = helpers.make_tuple(typechecker_api, list(column_types.values()))
return row_type
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 fail_if_manager_type_created_in_model_body(ctx: MethodContext) -> MypyType:
"""
Method hook that checks if method `<Manager>.from_queryset` is called inside a model class body.
Doing so won't, for instance, trigger the dynamic class hook(`create_new_manager_class_from_from_queryset_method`)
for managers.
"""
api = helpers.get_typechecker_api(ctx)
outer_model_info = api.scope.active_class()
if not outer_model_info or not outer_model_info.has_base(
fullnames.MODEL_CLASS_FULLNAME):
# Not inside a model class definition
return ctx.default_return_type
api.fail("`.from_queryset` called from inside model class body",
ctx.context,
code=errorcodes.MANAGER_UNTYPED)
return ctx.default_return_type
def transform_into_proper_return_type(
ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
default_return_type = ctx.default_return_type
assert isinstance(default_return_type, Instance)
outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
if outer_model_info is None or not helpers.is_model_subclass_info(
outer_model_info, django_context):
return ctx.default_return_type
assert isinstance(outer_model_info, TypeInfo)
if helpers.has_any_of_bases(default_return_type.type,
fullnames.RELATED_FIELDS_CLASSES):
return fill_descriptor_types_for_related_field(ctx, django_context)
if default_return_type.type.has_base(fullnames.ARRAY_FIELD_FULLNAME):
return determine_type_of_array_field(ctx, django_context)
return set_descriptor_types_for_field_callback(ctx, django_context)
def resolve_manager_method(ctx: AttributeContext) -> MypyType:
"""
A 'get_attribute_hook' that is intended to be invoked whenever the TypeChecker encounters
an attribute on a class that has 'django.db.models.BaseManager' as a base.
"""
api = helpers.get_typechecker_api(ctx)
# Skip (method) type that is currently something other than Any
if not isinstance(ctx.default_attr_type, AnyType):
return ctx.default_attr_type
# (Current state is:) We wouldn't end up here when looking up a method from a custom _manager_.
# That's why we only attempt to lookup the method for either a dynamically added or reverse manager.
assert isinstance(ctx.context, MemberExpr)
method_name = ctx.context.name
manager_instance = ctx.type
assert isinstance(manager_instance, Instance)
method_type = get_method_type_from_dynamic_manager(
api, method_name,
manager_instance.type) or get_method_type_from_reverse_manager(
api, method_name, manager_instance.type)
return method_type if method_type is not None else ctx.default_attr_type
def typecheck_model_method(ctx: Union[FunctionContext, MethodContext],
django_context: DjangoContext,
model_cls: Type[Model], method: str) -> MypyType:
typechecker_api = helpers.get_typechecker_api(ctx)
expected_types = django_context.get_expected_types(typechecker_api,
model_cls,
method=method)
expected_keys = [key for key in expected_types.keys() if key != 'pk']
for actual_name, actual_type in get_actual_types(ctx, expected_keys):
if actual_name not in expected_types:
ctx.api.fail(
'Unexpected attribute "{}" for model "{}"'.format(
actual_name, model_cls.__name__), ctx.context)
continue
helpers.check_types_compatible(
ctx,
expected_type=expected_types[actual_name],
actual_type=actual_type,
error_message='Incompatible type for "{}" of "{}"'.format(
actual_name, model_cls.__name__))
return ctx.default_return_type
def get_values_list_row_type(
ctx: MethodContext,
django_context: DjangoContext,
model_cls: Type[Model],
*,
is_annotated: bool,
flat: bool,
named: bool,
) -> MypyType:
field_lookups = resolve_field_lookups(ctx.args[0], django_context)
if field_lookups is None:
return AnyType(TypeOfAny.from_error)
typechecker_api = helpers.get_typechecker_api(ctx)
if len(field_lookups) == 0:
if flat:
primary_key_field = django_context.get_primary_key_field(model_cls)
lookup_type = get_field_type_from_lookup(
ctx,
django_context,
model_cls,
lookup=primary_key_field.attname,
method="values_list")
assert lookup_type is not None
return lookup_type
elif named:
column_types: "OrderedDict[str, MypyType]" = OrderedDict()
for field in django_context.get_model_fields(model_cls):
column_type = django_context.get_field_get_type(
typechecker_api, field, method="values_list")
column_types[field.attname] = column_type
if is_annotated:
# Return a NamedTuple with a fallback so that it's possible to access any field
return helpers.make_oneoff_named_tuple(
typechecker_api,
"Row",
column_types,
extra_bases=[
typechecker_api.named_generic_type(
ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])
],
)
else:
return helpers.make_oneoff_named_tuple(typechecker_api, "Row",
column_types)
else:
# flat=False, named=False, all fields
if is_annotated:
return typechecker_api.named_generic_type(
"builtins.tuple", [AnyType(TypeOfAny.special_form)])
field_lookups = []
for field in django_context.get_model_fields(model_cls):
field_lookups.append(field.attname)
if len(field_lookups) > 1 and flat:
typechecker_api.fail(
"'flat' is not valid when 'values_list' is called with more than one field",
ctx.context)
return AnyType(TypeOfAny.from_error)
column_types = OrderedDict()
for field_lookup in field_lookups:
lookup_field_type = get_field_type_from_lookup(
ctx,
django_context,
model_cls,
lookup=field_lookup,
method="values_list",
silent_on_error=is_annotated)
if lookup_field_type is None:
if is_annotated:
lookup_field_type = AnyType(TypeOfAny.from_omitted_generics)
else:
return AnyType(TypeOfAny.from_error)
column_types[field_lookup] = lookup_field_type
if flat:
assert len(column_types) == 1
row_type = next(iter(column_types.values()))
elif named:
row_type = helpers.make_oneoff_named_tuple(typechecker_api, "Row",
column_types)
else:
row_type = helpers.make_tuple(typechecker_api,
list(column_types.values()))
return row_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, [])
