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 return_proper_field_type_from_get_field(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# Options instance
assert isinstance(ctx.type, Instance)
model_type = ctx.type.args[0]
if not isinstance(model_type, Instance):
return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
model_cls = django_context.get_model_class_by_fullname(
model_type.type.fullname())
if model_cls is None:
return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
field_name_expr = helpers.get_call_argument_by_name(ctx, 'field_name')
if field_name_expr is None:
return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
field_name = helpers.resolve_string_attribute_value(
field_name_expr, ctx, django_context)
if field_name is None:
return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
try:
field = model_cls._meta.get_field(field_name)
except FieldDoesNotExist as exc:
ctx.api.fail(exc.args[0], ctx.context)
return AnyType(TypeOfAny.from_error)
field_fullname = helpers.get_class_fullname(field.__class__)
return _get_field_instance(ctx, field_fullname)
def extract_proper_type_queryset_values(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# called on QuerySet, return QuerySet of something
assert isinstance(ctx.type, Instance)
assert isinstance(ctx.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 ctx.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(ctx.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(ctx.default_return_type, [model_type, row_type])
def redefine_and_typecheck_model_init(
ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
assert isinstance(ctx.default_return_type, Instance)
model_fullname = ctx.default_return_type.type.fullname
model_cls = django_context.get_model_class_by_fullname(model_fullname)
if model_cls is None:
return ctx.default_return_type
return typecheck_model_method(ctx, django_context, model_cls, '__init__')
def redefine_and_typecheck_model_create(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
if not isinstance(ctx.default_return_type, Instance):
# only work with ctx.default_return_type = model Instance
return ctx.default_return_type
model_fullname = ctx.default_return_type.type.fullname
model_cls = django_context.get_model_class_by_fullname(model_fullname)
if model_cls is None:
return ctx.default_return_type
return typecheck_model_method(ctx, django_context, model_cls, 'create')
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 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
def return_proper_field_type_from_get_field(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# Options instance
assert isinstance(ctx.type, Instance)
# bail if list of generic params is empty
if len(ctx.type.args) == 0:
return ctx.default_return_type
model_type = ctx.type.args[0]
if not isinstance(model_type, Instance):
return ctx.default_return_type
model_cls = django_context.get_model_class_by_fullname(
model_type.type.fullname)
if model_cls is None:
return ctx.default_return_type
field_name_expr = helpers.get_call_argument_by_name(ctx, 'field_name')
if field_name_expr is None:
return ctx.default_return_type
field_name = helpers.resolve_string_attribute_value(
field_name_expr, django_context)
if field_name is None:
return ctx.default_return_type
try:
field = model_cls._meta.get_field(field_name)
except FieldDoesNotExist as exc:
# if model is abstract, do not raise exception, skip false positives
if not model_cls._meta.abstract:
ctx.api.fail(exc.args[0], ctx.context)
return AnyType(TypeOfAny.from_error)
field_fullname = helpers.get_class_fullname(field.__class__)
return _get_field_instance(ctx, field_fullname)
