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_model_pk_lookup(api: CheckerPluginInterface,
model_type_info: TypeInfo) -> LookupNode:
# Primary keys are special-cased
primary_key_type = helpers.extract_primary_key_type_for_get(
model_type_info)
if primary_key_type:
return FieldNode(primary_key_type)
else:
# No PK, use the get type for AutoField as PK type.
autofield_info = api.lookup_typeinfo(
'django.db.models.fields.AutoField')
pk_type = helpers.get_private_descriptor_type(autofield_info,
'_pyi_private_get_type',
is_nullable=False)
return FieldNode(pk_type)
def resolve_model_lookup(api: CheckerPluginInterface,
model_type_info: TypeInfo, lookup: str) -> LookupNode:
"""Resolve a lookup on the given model."""
if lookup == 'pk':
# Primary keys are special-cased
primary_key_type = helpers.extract_primary_key_type_for_get(
model_type_info)
if primary_key_type:
return FieldNode(primary_key_type)
else:
# No PK, use the get type for AutoField as PK type.
autofield_info = api.lookup_typeinfo(
'django.db.models.fields.AutoField')
pk_type = helpers.get_private_descriptor_type(
autofield_info, '_pyi_private_get_type', is_nullable=False)
return FieldNode(pk_type)
field_name = get_actual_field_name_for_lookup_field(
lookup, model_type_info)
field_node = model_type_info.get(field_name)
if not field_node:
raise LookupException(
f'When resolving lookup "{lookup}", field "{field_name}" was not found in model {model_type_info.name()}'
)
if field_name.endswith('_id'):
field_name_without_id = field_name.rstrip('_id')
foreign_key_field = model_type_info.get(field_name_without_id)
if foreign_key_field is not None and helpers.is_foreign_key(
foreign_key_field.type):
# Hack: If field ends with '_id' and there is a model field without the '_id' suffix, then use that field.
field_node = foreign_key_field
field_name = field_name_without_id
field_node_type = field_node.type
if field_node_type is None or not isinstance(field_node_type, Instance):
raise LookupException(
f'When resolving lookup "{lookup}", could not determine type for {model_type_info.name()}.{field_name}'
)
if helpers.is_foreign_key(field_node_type):
field_type = helpers.extract_field_getter_type(field_node_type)
is_nullable = helpers.is_optional(field_type)
if is_nullable:
field_type = helpers.make_required(field_type)
if isinstance(field_type, Instance):
return RelatedModelNode(typ=field_type, is_nullable=is_nullable)
else:
raise LookupException(
f"Not an instance for field {field_type} lookup {lookup}")
field_type = helpers.extract_field_getter_type(field_node_type)
if field_type:
return FieldNode(typ=field_type)
else:
# Not a Field
if field_name == 'id':
# If no 'id' field was fouond, use an int
return FieldNode(api.named_generic_type('builtins.int', []))
related_manager_arg = None
if field_node_type.type.has_base(
helpers.RELATED_MANAGER_CLASS_FULLNAME):
related_manager_arg = field_node_type.args[0]
if related_manager_arg is not None:
# Reverse relation
return RelatedModelNode(typ=related_manager_arg, is_nullable=True)
raise LookupException(
f'When resolving lookup "{lookup}", could not determine type for {model_type_info.name()}.{field_name}'
)
