def extract_proper_type_for_values_list(ctx: MethodContext) -> Type:
object_type = ctx.type
if not isinstance(object_type, Instance):
return ctx.default_return_type
flat = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'flat'))
named = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'named'))
ret = ctx.default_return_type
any_type = AnyType(TypeOfAny.implementation_artifact)
if named and flat:
ctx.api.fail("'flat' and 'named' can't be used together.", ctx.context)
return ret
elif named:
# TODO: Fill in namedtuple fields/types
row_arg = ctx.api.named_generic_type('typing.NamedTuple', [])
elif flat:
# TODO: Figure out row_arg type dependent on the argument passed in
if len(ctx.args[0]) > 1:
ctx.api.fail(
"'flat' is not valid when values_list is called with more than one field.",
ctx.context)
return ret
row_arg = any_type
else:
# TODO: Figure out tuple argument types dependent on the arguments passed in
row_arg = ctx.api.named_generic_type('builtins.tuple', [any_type])
first_arg = ret.args[0] if len(ret.args) > 0 else any_type
new_type_args = [first_arg, row_arg]
return helpers.reparametrize_instance(ret, new_type_args)
def determine_proper_manager_type(ctx: FunctionContext) -> Type:
from mypy.checker import TypeChecker
api = cast(TypeChecker, ctx.api)
ret = ctx.default_return_type
if not api.tscope.classes:
# not in class
return ret
outer_model_info = api.tscope.classes[0]
if not outer_model_info.has_base(helpers.MODEL_CLASS_FULLNAME):
return ret
if not isinstance(ret, Instance):
return ret
has_manager_base = False
for i, base in enumerate(ret.type.bases):
if base.type.fullname() in {
helpers.MANAGER_CLASS_FULLNAME,
helpers.RELATED_MANAGER_CLASS_FULLNAME,
helpers.BASE_MANAGER_CLASS_FULLNAME
}:
has_manager_base = True
break
if has_manager_base:
# Fill in the manager's type argument from the outer model
new_type_args = [Instance(outer_model_info, [])]
return helpers.reparametrize_instance(ret, new_type_args)
else:
return ret
def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance:
default_return_type = cast(Instance, ctx.default_return_type)
is_nullable = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'null'))
set_type = helpers.get_private_descriptor_type(default_return_type.type,
'_pyi_private_set_type',
is_nullable=is_nullable)
get_type = helpers.get_private_descriptor_type(default_return_type.type,
'_pyi_private_get_type',
is_nullable=is_nullable)
return helpers.reparametrize_instance(default_return_type,
[set_type, get_type])
def fill_descriptor_types_for_related_field(ctx: FunctionContext) -> Type:
default_return_type = set_descriptor_types_for_field(ctx)
referred_to_type = _extract_referred_to_type(ctx)
if referred_to_type is None:
return default_return_type
# replace Any with referred_to_type
args = []
for default_arg in default_return_type.args:
args.append(convert_any_to_type(default_arg, referred_to_type))
return helpers.reparametrize_instance(ctx.default_return_type, new_args=args)
def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance:
default_return_type = cast(Instance, ctx.default_return_type)
is_nullable = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'null'))
if not is_nullable and default_return_type.type.has_base(helpers.CHAR_FIELD_FULLNAME):
# blank=True for CharField can be interpreted as null=True
is_nullable = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'blank'))
set_type = get_private_descriptor_type(default_return_type.type, '_pyi_private_set_type',
is_nullable=is_nullable)
get_type = get_private_descriptor_type(default_return_type.type, '_pyi_private_get_type',
is_nullable=is_nullable)
return helpers.reparametrize_instance(default_return_type, [set_type, get_type])
def determine_type_of_array_field(ctx: FunctionContext) -> Type:
default_return_type = set_descriptor_types_for_field(ctx)
base_field_arg_type = helpers.get_argument_type_by_name(ctx, 'base_field')
if not base_field_arg_type or not isinstance(base_field_arg_type, Instance):
return default_return_type
base_type = base_field_arg_type.args[1] # extract __get__ type
args = []
for default_arg in default_return_type.args:
args.append(convert_any_to_type(default_arg, base_type))
return helpers.reparametrize_instance(default_return_type, args)
def extract_proper_type_for_queryset_values(ctx: MethodContext) -> Type:
object_type = ctx.type
if not isinstance(object_type, Instance):
return ctx.default_return_type
fields_arg_expr = ctx.args[ctx.callee_arg_names.index('fields')]
if len(fields_arg_expr) == 0:
# values_list/values with no args is not yet supported, so default to Any types for field types
# It should in the future include all model fields, "extra" fields and "annotated" fields
return ctx.default_return_type
model_arg = get_queryset_model_arg(ctx.default_return_type)
if isinstance(model_arg, Instance):
model_type_info = model_arg.type
else:
model_type_info = None
column_types: OrderedDict[str, Type] = OrderedDict()
# parse *fields
for field_expr in fields_arg_expr:
if isinstance(field_expr, StrExpr):
field_name = field_expr.value
# Default to any type
column_types[field_name] = AnyType(
TypeOfAny.implementation_artifact)
if model_type_info:
resolved_lookup_type = resolve_values_lookup(
ctx.api, model_type_info, field_name)
if resolved_lookup_type is not None:
column_types[field_name] = resolved_lookup_type
else:
return ctx.default_return_type
# parse **expressions
expression_arg_names = ctx.arg_names[ctx.callee_arg_names.index(
'expressions')]
for expression_name in expression_arg_names:
# Arbitrary additional annotation expressions are supported, but they all have type Any for now
column_types[expression_name] = AnyType(
TypeOfAny.implementation_artifact)
row_arg = helpers.make_typeddict(ctx.api,
fields=column_types,
required_keys=set())
return helpers.reparametrize_instance(ctx.default_return_type,
[model_arg, row_arg])
def convert_any_to_type(typ: Type, referred_to_type: Type) -> Type:
if isinstance(typ, UnionType):
converted_items = []
for item in typ.items:
converted_items.append(convert_any_to_type(item, referred_to_type))
return UnionType.make_simplified_union(converted_items,
line=typ.line, column=typ.column)
if isinstance(typ, Instance):
args = []
for default_arg in typ.args:
if isinstance(default_arg, AnyType):
args.append(referred_to_type)
else:
args.append(default_arg)
return helpers.reparametrize_instance(typ, args)
if isinstance(typ, AnyType):
return referred_to_type
return typ
def extract_proper_type_queryset_values_list(ctx: MethodContext) -> Type:
object_type = ctx.type
if not isinstance(object_type, Instance):
return ctx.default_return_type
ret = ctx.default_return_type
model_arg = get_queryset_model_arg(ctx.default_return_type)
# model_arg: Union[AnyType, Type] = ret.args[0] if len(ret.args) > 0 else any_type
column_names: List[Optional[str]] = []
column_types: OrderedDict[str, Type] = OrderedDict()
fields_arg_expr = ctx.args[ctx.callee_arg_names.index('fields')]
fields_param_is_specified = True
if len(fields_arg_expr) == 0:
# values_list/values with no args is not yet supported, so default to Any types for field types
# It should in the future include all model fields, "extra" fields and "annotated" fields
fields_param_is_specified = False
if isinstance(model_arg, Instance):
model_type_info = model_arg.type
else:
model_type_info = None
any_type = AnyType(TypeOfAny.implementation_artifact)
# Figure out each field name passed to fields
only_strings_as_fields_expressions = True
for field_expr in fields_arg_expr:
if isinstance(field_expr, StrExpr):
field_name = field_expr.value
column_names.append(field_name)
# Default to any type
column_types[field_name] = any_type
if model_type_info:
resolved_lookup_type = resolve_values_lookup(
ctx.api, model_type_info, field_name)
if resolved_lookup_type is not None:
column_types[field_name] = resolved_lookup_type
else:
# Dynamic field names are partially supported for values_list, but not values
column_names.append(None)
only_strings_as_fields_expressions = False
flat = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'flat'))
named = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'named'))
api = cast(TypeChecker, ctx.api)
if named and flat:
api.fail("'flat' and 'named' can't be used together.", ctx.context)
return ret
elif named:
# named=True, flat=False -> List[NamedTuple]
if fields_param_is_specified and only_strings_as_fields_expressions:
row_arg = helpers.make_named_tuple(api,
fields=column_types,
name="Row")
else:
# fallback to catch-all NamedTuple
row_arg = helpers.make_named_tuple(api,
fields=OrderedDict(),
name="Row")
elif flat:
# named=False, flat=True -> List of elements
if len(ctx.args[0]) > 1:
api.fail(
"'flat' is not valid when values_list is called with more than one field.",
ctx.context)
return ctx.default_return_type
if fields_param_is_specified and only_strings_as_fields_expressions:
# Grab first element
row_arg = column_types[column_names[0]]
else:
row_arg = any_type
else:
# named=False, flat=False -> List[Tuple]
if fields_param_is_specified:
args = [
# Fallback to Any if the column name is unknown (e.g. dynamic)
column_types.get(column_name, any_type)
if column_name is not None else any_type
for column_name in column_names
]
else:
args = [any_type]
row_arg = helpers.make_tuple(api, fields=args)
new_type_args = [model_arg, row_arg]
return helpers.reparametrize_instance(ret, new_type_args)
def extract_proper_type_for_values_and_values_list(method_name: str,
ctx: MethodContext) -> Type:
api = cast(TypeChecker, ctx.api)
object_type = ctx.type
if not isinstance(object_type, Instance):
return ctx.default_return_type
ret = ctx.default_return_type
any_type = AnyType(TypeOfAny.implementation_artifact)
fields_arg_expr = ctx.args[ctx.callee_arg_names.index('fields')]
model_arg: Union[AnyType,
Type] = ret.args[0] if len(ret.args) > 0 else any_type
column_names: List[Optional[str]] = []
column_types: OrderedDict[str, Type] = OrderedDict()
fill_column_types = True
if len(fields_arg_expr) == 0:
# values_list/values with no args is not yet supported, so default to Any types for field types
# It should in the future include all model fields, "extra" fields and "annotated" fields
fill_column_types = False
if isinstance(model_arg, Instance):
model_type_info = model_arg.type
else:
model_type_info = None
# Figure out each field name passed to fields
has_dynamic_column_names = False
for field_expr in fields_arg_expr:
if isinstance(field_expr, StrExpr):
field_name = field_expr.value
column_names.append(field_name)
# Default to any type
column_types[field_name] = any_type
if model_type_info:
resolved_lookup_type = resolve_values_lookup(
ctx.api, model_type_info, field_name)
if resolved_lookup_type is not None:
column_types[field_name] = resolved_lookup_type
else:
# Dynamic field names are partially supported for values_list, but not values
column_names.append(None)
has_dynamic_column_names = True
if method_name == 'values_list':
flat = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'flat'))
named = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'named'))
if named and flat:
api.fail("'flat' and 'named' can't be used together.", ctx.context)
return ret
elif named:
if fill_column_types and not has_dynamic_column_names:
row_arg = helpers.make_named_tuple(api,
fields=column_types,
name="Row")
else:
row_arg = helpers.make_named_tuple(api,
fields=OrderedDict(),
name="Row")
elif flat:
if len(ctx.args[0]) > 1:
api.fail(
"'flat' is not valid when values_list is called with more than one field.",
ctx.context)
return ret
if fill_column_types and not has_dynamic_column_names:
# Grab first element
row_arg = column_types[column_names[0]]
else:
row_arg = any_type
else:
if fill_column_types:
args = [
# Fallback to Any if the column name is unknown (e.g. dynamic)
column_types.get(column_name, any_type)
if column_name is not None else any_type
for column_name in column_names
]
else:
args = [any_type]
row_arg = helpers.make_tuple(api, fields=args)
elif method_name == 'values':
expression_arg_names = ctx.arg_names[ctx.callee_arg_names.index(
'expressions')]
for expression_name in expression_arg_names:
# Arbitrary additional annotation expressions are supported, but they all have type Any for now
column_names.append(expression_name)
column_types[expression_name] = any_type
if fill_column_types and not has_dynamic_column_names:
row_arg = helpers.make_typeddict(api,
fields=column_types,
required_keys=set())
else:
return ctx.default_return_type
else:
raise Exception(
f"extract_proper_type_for_values_list doesn't support method {method_name}"
)
new_type_args = [model_arg, row_arg]
return helpers.reparametrize_instance(ret, new_type_args)
