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_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 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
