def extract_and_return_primary_key_of_bound_related_field_parameter(ctx: AttributeContext) -> Type:
if not isinstance(ctx.default_attr_type, Instance) or not (ctx.default_attr_type.type.fullname() == 'builtins.int'):
return ctx.default_attr_type
if not isinstance(ctx.type, Instance) or not ctx.type.type.has_base(helpers.MODEL_CLASS_FULLNAME):
return ctx.default_attr_type
field_name = ctx.context.name.split('_')[0]
sym = ctx.type.type.get(field_name)
if sym and isinstance(sym.type, Instance) and len(sym.type.args) > 0:
referred_to = sym.type.args[1]
if isinstance(referred_to, AnyType):
return AnyType(TypeOfAny.implementation_artifact)
model_type = _extract_referred_to_type_info(referred_to)
if model_type is None:
return AnyType(TypeOfAny.implementation_artifact)
primary_key_type = helpers.extract_primary_key_type_for_get(model_type)
if primary_key_type:
return primary_key_type
is_nullable = helpers.get_fields_metadata(ctx.type.type).get(field_name, {}).get('null', False)
if is_nullable:
return helpers.make_optional(ctx.default_attr_type)
return ctx.default_attr_type
def resolve_values_lookup(api: CheckerPluginInterface,
model_type_info: TypeInfo,
lookup: str) -> Optional[Type]:
"""Resolves a values/values_list lookup if possible, to a Type."""
try:
nodes = resolve_lookup(api, model_type_info, lookup)
except LookupException:
nodes = []
if not nodes:
return None
make_optional = False
for node in nodes:
if isinstance(node, RelatedModelNode) and node.is_nullable:
# All lookups following a relation which is nullable should be optional
make_optional = True
node = nodes[-1]
node_type = node.typ
if isinstance(node, RelatedModelNode):
# Related models used in values/values_list get resolved to the primary key of the related model.
# So, we lookup the pk of that model.
pk_lookup_nodes = resolve_lookup(api, node_type.type, "pk")
if not pk_lookup_nodes:
return None
node_type = pk_lookup_nodes[0].typ
if make_optional:
return helpers.make_optional(node_type)
else:
return node_type
def get_private_descriptor_type(type_info: TypeInfo, private_field_name: str, is_nullable: bool) -> Type:
node = type_info.get(private_field_name).node
if isinstance(node, Var):
descriptor_type = node.type
if is_nullable:
descriptor_type = helpers.make_optional(descriptor_type)
return descriptor_type
return AnyType(TypeOfAny.unannotated)
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