def get_or_create_annotated_type(
api: Union[SemanticAnalyzer,
CheckerPluginInterface], model_type: Instance,
fields_dict: Optional[TypedDictType]) -> Instance:
"""
Get or create the type for a model for which you getting/setting any attr is allowed.
The generated type is an subclass of the model and django._AnyAttrAllowed.
The generated type is placed in the django_stubs_ext module, with the name WithAnnotations[ModelName].
If the user wanted to annotate their code using this type, then this is the annotation they would use.
This is a bit of a hack to make a pretty type for error messages and which would make sense for users.
"""
model_module_name = "django_stubs_ext"
if helpers.is_annotated_model_fullname(model_type.type.fullname):
# If it's already a generated class, we want to use the original model as a base
model_type = model_type.type.bases[0]
if fields_dict is not None:
type_name = f"WithAnnotations[{model_type.type.fullname.replace('.', '__')}, {fields_dict}]"
else:
type_name = f"WithAnnotations[{model_type.type.fullname.replace('.', '__')}]"
annotated_typeinfo = helpers.lookup_fully_qualified_typeinfo(
cast(TypeChecker, api), model_module_name + "." + type_name)
if annotated_typeinfo is None:
model_module_file = api.modules[model_module_name] # type: ignore
if isinstance(api, SemanticAnalyzer):
annotated_model_type = api.named_type_or_none(
ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])
assert annotated_model_type is not None
else:
annotated_model_type = api.named_generic_type(
ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])
annotated_typeinfo = add_new_class_for_module(
model_module_file,
type_name,
bases=[model_type]
if fields_dict is not None else [model_type, annotated_model_type],
fields=fields_dict.items if fields_dict is not None else None,
no_serialize=True,
)
if fields_dict is not None:
# To allow structural subtyping, make it a Protocol
annotated_typeinfo.is_protocol = True
# Save for later to easily find which field types were annotated
annotated_typeinfo.metadata[
"annotated_field_types"] = fields_dict.items
annotated_type = Instance(annotated_typeinfo, [])
return annotated_type
def typecheck_queryset_filter(ctx: MethodContext,
django_context: DjangoContext) -> MypyType:
# Expected formal arguments for filter methods are `*args` and `**kwargs`. We'll only typecheck
# `**kwargs`, which means that `arg_names[1]` is what we're interested in.
lookup_kwargs = ctx.arg_names[1]
provided_lookup_types = ctx.arg_types[1]
assert isinstance(ctx.type, Instance)
if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
return ctx.default_return_type
model_cls_fullname = ctx.type.args[0].type.fullname
model_cls = django_context.get_model_class_by_fullname(model_cls_fullname)
if model_cls is None:
return ctx.default_return_type
for lookup_kwarg, provided_type in zip(lookup_kwargs,
provided_lookup_types):
if lookup_kwarg is None:
continue
if isinstance(provided_type, Instance) and provided_type.type.has_base(
fullnames.COMBINABLE_EXPRESSION_FULLNAME):
provided_type = resolve_combinable_type(provided_type,
django_context)
lookup_type: MypyType
if is_annotated_model_fullname(model_cls_fullname):
lookup_type = AnyType(TypeOfAny.implementation_artifact)
else:
lookup_type = django_context.resolve_lookup_expected_type(
ctx, model_cls, lookup_kwarg)
# Managers as provided_type is not supported yet
if isinstance(provided_type, Instance) and helpers.has_any_of_bases(
provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME,
fullnames.QUERYSET_CLASS_FULLNAME)):
return ctx.default_return_type
helpers.check_types_compatible(
ctx,
expected_type=lookup_type,
actual_type=provided_type,
error_message=f"Incompatible type for lookup {lookup_kwarg!r}:",
)
return ctx.default_return_type
def extract_proper_type_queryset_values_list(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# called on the Instance, returns QuerySet of something
assert isinstance(ctx.type, Instance)
default_return_type = get_proper_type(ctx.default_return_type)
assert isinstance(default_return_type, Instance)
model_type = _extract_model_type_from_queryset(ctx.type)
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)
model_cls = django_context.get_model_class_by_fullname(
model_type.type.fullname)
if model_cls is None:
return default_return_type
flat_expr = helpers.get_call_argument_by_name(ctx, "flat")
if flat_expr is not None and isinstance(flat_expr, NameExpr):
flat = helpers.parse_bool(flat_expr)
else:
flat = False
named_expr = helpers.get_call_argument_by_name(ctx, "named")
if named_expr is not None and isinstance(named_expr, NameExpr):
named = helpers.parse_bool(named_expr)
else:
named = False
if flat and named:
ctx.api.fail("'flat' and 'named' can't be used together", ctx.context)
return helpers.reparametrize_instance(
default_return_type,
[model_type, AnyType(TypeOfAny.from_error)])
# account for possible None
flat = flat or False
named = named or False
is_annotated = is_annotated_model_fullname(model_type.type.fullname)
row_type = get_values_list_row_type(ctx,
django_context,
model_cls,
is_annotated=is_annotated,
flat=flat,
named=named)
return helpers.reparametrize_instance(default_return_type,
[model_type, row_type])
def extract_proper_type_queryset_values(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# called on QuerySet, return QuerySet of something
assert isinstance(ctx.type, Instance)
default_return_type = get_proper_type(ctx.default_return_type)
assert isinstance(default_return_type, Instance)
model_type = _extract_model_type_from_queryset(ctx.type)
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)
model_cls = django_context.get_model_class_by_fullname(
model_type.type.fullname)
if model_cls is None:
return default_return_type
if is_annotated_model_fullname(model_type.type.fullname):
return default_return_type
field_lookups = resolve_field_lookups(ctx.args[0], django_context)
if field_lookups is None:
return AnyType(TypeOfAny.from_error)
if len(field_lookups) == 0:
for field in django_context.get_model_fields(model_cls):
field_lookups.append(field.attname)
column_types: "OrderedDict[str, MypyType]" = OrderedDict()
for field_lookup in field_lookups:
field_lookup_type = get_field_type_from_lookup(ctx,
django_context,
model_cls,
lookup=field_lookup,
method="values")
if field_lookup_type is None:
return helpers.reparametrize_instance(
default_return_type,
[model_type, AnyType(TypeOfAny.from_error)])
column_types[field_lookup] = field_lookup_type
row_type = helpers.make_typeddict(ctx.api, column_types,
set(column_types.keys()))
return helpers.reparametrize_instance(default_return_type,
[model_type, row_type])
def typecheck_queryset_filter(ctx: MethodContext,
django_context: DjangoContext) -> MypyType:
lookup_kwargs = ctx.arg_names[1]
provided_lookup_types = ctx.arg_types[1]
assert isinstance(ctx.type, Instance)
if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
return ctx.default_return_type
model_cls_fullname = ctx.type.args[0].type.fullname
model_cls = django_context.get_model_class_by_fullname(model_cls_fullname)
if model_cls is None:
return ctx.default_return_type
for lookup_kwarg, provided_type in zip(lookup_kwargs,
provided_lookup_types):
if lookup_kwarg is None:
continue
if isinstance(provided_type, Instance) and provided_type.type.has_base(
"django.db.models.expressions.Combinable"):
provided_type = resolve_combinable_type(provided_type,
django_context)
lookup_type: MypyType
if is_annotated_model_fullname(model_cls_fullname):
lookup_type = AnyType(TypeOfAny.implementation_artifact)
else:
lookup_type = django_context.resolve_lookup_expected_type(
ctx, model_cls, lookup_kwarg)
# Managers as provided_type is not supported yet
if isinstance(provided_type, Instance) and helpers.has_any_of_bases(
provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME,
fullnames.QUERYSET_CLASS_FULLNAME)):
return ctx.default_return_type
helpers.check_types_compatible(
ctx,
expected_type=lookup_type,
actual_type=provided_type,
error_message=f"Incompatible type for lookup {lookup_kwarg!r}:",
)
return ctx.default_return_type