def transform_into_proper_return_type(ctx: FunctionContext) -> Type:
default_return_type = ctx.default_return_type
if not isinstance(default_return_type, Instance):
return default_return_type
if helpers.has_any_of_bases(default_return_type.type, (helpers.FOREIGN_KEY_FULLNAME,
helpers.ONETOONE_FIELD_FULLNAME,
helpers.MANYTOMANY_FIELD_FULLNAME)):
return fill_descriptor_types_for_related_field(ctx)
if default_return_type.type.has_base(helpers.ARRAY_FIELD_FULLNAME):
return determine_type_of_array_field(ctx)
return set_descriptor_types_for_field(ctx)
def extract_expected_types(ctx: FunctionContext,
model: TypeInfo,
is_init: bool = False) -> Dict[str, Type]:
api = cast(TypeChecker, ctx.api)
expected_types: Dict[str, Type] = {}
primary_key_type = helpers.extract_explicit_set_type_of_model_primary_key(
model)
if not primary_key_type:
# no explicit primary key, set pk to Any and add id
primary_key_type = AnyType(TypeOfAny.special_form)
if is_init:
expected_types['id'] = helpers.make_optional(
ctx.api.named_generic_type('builtins.int', []))
else:
expected_types['id'] = ctx.api.named_generic_type(
'builtins.int', [])
expected_types['pk'] = primary_key_type
for base in model.mro:
# extract all fields for all models in MRO
for name, sym in base.names.items():
# do not redefine special attrs
if name in {'_meta', 'pk'}:
continue
if isinstance(sym.node, Var):
typ = sym.node.type
if typ is None or isinstance(typ, AnyType):
# types are not ready, fallback to Any
expected_types[name] = AnyType(
TypeOfAny.from_unimported_type)
expected_types[name + '_id'] = AnyType(
TypeOfAny.from_unimported_type)
elif isinstance(typ, Instance):
field_type = helpers.extract_field_setter_type(typ)
if field_type is None:
continue
if helpers.has_any_of_bases(
typ.type, (helpers.FOREIGN_KEY_FULLNAME,
helpers.ONETOONE_FIELD_FULLNAME)):
related_primary_key_type = AnyType(
TypeOfAny.implementation_artifact)
# in case it's optional, we need Instance type
referred_to_model = typ.args[1]
is_nullable = helpers.is_optional(referred_to_model)
if is_nullable:
referred_to_model = helpers.make_required(
typ.args[1])
if isinstance(
referred_to_model,
Instance) and referred_to_model.type.has_base(
helpers.MODEL_CLASS_FULLNAME):
pk_type = helpers.extract_explicit_set_type_of_model_primary_key(
referred_to_model.type)
if not pk_type:
# extract set type of AutoField
autofield_info = api.lookup_typeinfo(
'django.db.models.fields.AutoField')
pk_type = helpers.get_private_descriptor_type(
autofield_info,
'_pyi_private_set_type',
is_nullable=is_nullable)
related_primary_key_type = pk_type
if is_init:
related_primary_key_type = helpers.make_optional(
related_primary_key_type)
expected_types[name + '_id'] = related_primary_key_type
field_metadata = helpers.get_fields_metadata(model).get(
name, {})
if field_type:
# related fields could be None in __init__ (but should be specified before save())
if helpers.has_any_of_bases(
typ.type,
(helpers.FOREIGN_KEY_FULLNAME,
helpers.ONETOONE_FIELD_FULLNAME)) and is_init:
field_type = helpers.make_optional(field_type)
# if primary_key=True and default specified
elif field_metadata.get(
'primary_key', False) and field_metadata.get(
'default_specified', False):
field_type = helpers.make_optional(field_type)
# if CharField(blank=True,...) and not nullable, then field can be None in __init__
elif (helpers.has_any_of_bases(
typ.type,
(helpers.CHAR_FIELD_FULLNAME, )) and is_init
and field_metadata.get('blank', False)
and not field_metadata.get('null', False)):
field_type = helpers.make_optional(field_type)
expected_types[name] = field_type
return expected_types