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 extract_proper_type_queryset_annotate(
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)
api = ctx.api
field_types = model_type.type.metadata.get("annotated_field_types")
kwargs = gather_kwargs(ctx)
if kwargs:
# For now, we don't try to resolve the output_field of the field would be, but use Any.
added_field_types = {
name: AnyType(TypeOfAny.implementation_artifact)
for name, typ in kwargs.items()
}
if field_types is not None:
# Annotate was called more than once, so add/update existing field types
field_types.update(added_field_types)
else:
field_types = added_field_types
fields_dict = None
if field_types is not None:
fields_dict = helpers.make_typeddict(api,
fields=OrderedDict(field_types),
required_keys=set(
field_types.keys()))
annotated_type = get_or_create_annotated_type(api,
model_type,
fields_dict=fields_dict)
row_type: MypyType
if len(default_return_type.args) > 1:
original_row_type: MypyType = default_return_type.args[1]
row_type = original_row_type
if isinstance(original_row_type, TypedDictType):
row_type = api.named_generic_type("builtins.dict", [
api.named_generic_type("builtins.str", []),
AnyType(TypeOfAny.from_omitted_generics)
])
elif isinstance(original_row_type, TupleType):
fallback: Instance = original_row_type.partial_fallback
if fallback is not None and fallback.type.has_base(
"typing.NamedTuple"):
# TODO: Use a NamedTuple which contains the known fields, but also
# falls back to allowing any attribute access.
row_type = AnyType(TypeOfAny.implementation_artifact)
else:
row_type = api.named_generic_type(
"builtins.tuple",
[AnyType(TypeOfAny.from_omitted_generics)])
elif isinstance(original_row_type,
Instance) and original_row_type.type.has_base(
fullnames.MODEL_CLASS_FULLNAME):
row_type = annotated_type
else:
row_type = annotated_type
return helpers.reparametrize_instance(default_return_type,
[annotated_type, row_type])
