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}' )