[ty] Implemented partial support for "find references" language server feature. (#19475)

This PR adds basic support for the "find all references" language server feature.

---------

Co-authored-by: UnboundVariable <unbound@gmail.com>
This commit is contained in:
UnboundVariable 2025-07-23 09:16:22 -07:00 committed by GitHub
parent 89258f1938
commit fa1df4cedc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1468 additions and 169 deletions

View file

@ -235,12 +235,7 @@ impl TraversalSignal {
}
pub fn walk_annotation<'a, V: SourceOrderVisitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
let node = AnyNodeRef::from(expr);
if visitor.enter_node(node).is_traverse() {
visitor.visit_expr(expr);
}
visitor.leave_node(node);
}
pub fn walk_decorator<'a, V>(visitor: &mut V, decorator: &'a Decorator)

View file

@ -52,9 +52,7 @@ pub(crate) fn covering_node(root: AnyNodeRef, range: TextRange) -> CoveringNode
if visitor.ancestors.is_empty() {
visitor.ancestors.push(root);
}
CoveringNode {
nodes: visitor.ancestors,
}
CoveringNode::from_ancestors(visitor.ancestors)
}
/// The node with a minimal range that fully contains the search range.
@ -67,6 +65,12 @@ pub(crate) struct CoveringNode<'a> {
}
impl<'a> CoveringNode<'a> {
/// Creates a new `CoveringNode` from a list of ancestor nodes.
/// The ancestors should be ordered from root to the covering node.
pub(crate) fn from_ancestors(ancestors: Vec<AnyNodeRef<'a>>) -> Self {
Self { nodes: ancestors }
}
/// Returns the covering node found.
pub(crate) fn node(&self) -> AnyNodeRef<'a> {
*self

View file

@ -2,6 +2,8 @@ pub use crate::goto_declaration::goto_declaration;
pub use crate::goto_definition::goto_definition;
pub use crate::goto_type_definition::goto_type_definition;
use std::borrow::Cow;
use crate::find_node::covering_node;
use crate::stub_mapping::StubMapper;
use ruff_db::parsed::ParsedModuleRef;
@ -270,10 +272,270 @@ impl GotoTarget<'_> {
definitions_to_navigation_targets(db, stub_mapper, definitions)
}
// For exception variables, they are their own definitions (like parameters)
GotoTarget::ExceptVariable(except_handler) => {
if let Some(name) = &except_handler.name {
let range = name.range;
Some(crate::NavigationTargets::single(NavigationTarget::new(
file, range,
)))
} else {
None
}
}
// For pattern match rest variables, they are their own definitions
GotoTarget::PatternMatchRest(pattern_mapping) => {
if let Some(rest_name) = &pattern_mapping.rest {
let range = rest_name.range;
Some(crate::NavigationTargets::single(NavigationTarget::new(
file, range,
)))
} else {
None
}
}
// For pattern match as names, they are their own definitions
GotoTarget::PatternMatchAsName(pattern_as) => {
if let Some(name) = &pattern_as.name {
let range = name.range;
Some(crate::NavigationTargets::single(NavigationTarget::new(
file, range,
)))
} else {
None
}
}
// TODO: Handle string literals that map to TypedDict fields
_ => None,
}
}
/// Returns the text representation of this goto target.
/// Returns `None` if no meaningful string representation can be provided.
/// This is used by the "references" feature, which looks for references
/// to this goto target.
pub(crate) fn to_string(&self) -> Option<Cow<str>> {
match self {
GotoTarget::Expression(expression) => match expression {
ast::ExprRef::Name(name) => Some(Cow::Borrowed(name.id.as_str())),
ast::ExprRef::Attribute(attr) => Some(Cow::Borrowed(attr.attr.as_str())),
_ => None,
},
GotoTarget::FunctionDef(function) => Some(Cow::Borrowed(function.name.as_str())),
GotoTarget::ClassDef(class) => Some(Cow::Borrowed(class.name.as_str())),
GotoTarget::Parameter(parameter) => Some(Cow::Borrowed(parameter.name.as_str())),
GotoTarget::ImportSymbolAlias { alias, .. } => {
if let Some(asname) = &alias.asname {
Some(Cow::Borrowed(asname.as_str()))
} else {
Some(Cow::Borrowed(alias.name.as_str()))
}
}
GotoTarget::ImportModuleComponent {
module_name,
component_index,
..
} => {
let components: Vec<&str> = module_name.split('.').collect();
if let Some(component) = components.get(*component_index) {
Some(Cow::Borrowed(*component))
} else {
Some(Cow::Borrowed(module_name))
}
}
GotoTarget::ImportModuleAlias { alias } => {
if let Some(asname) = &alias.asname {
Some(Cow::Borrowed(asname.as_str()))
} else {
Some(Cow::Borrowed(alias.name.as_str()))
}
}
GotoTarget::ExceptVariable(except) => {
Some(Cow::Borrowed(except.name.as_ref()?.as_str()))
}
GotoTarget::KeywordArgument { keyword, .. } => {
Some(Cow::Borrowed(keyword.arg.as_ref()?.as_str()))
}
GotoTarget::PatternMatchRest(rest) => Some(Cow::Borrowed(rest.rest.as_ref()?.as_str())),
GotoTarget::PatternKeywordArgument(keyword) => {
Some(Cow::Borrowed(keyword.attr.as_str()))
}
GotoTarget::PatternMatchStarName(star) => {
Some(Cow::Borrowed(star.name.as_ref()?.as_str()))
}
GotoTarget::PatternMatchAsName(as_name) => {
Some(Cow::Borrowed(as_name.name.as_ref()?.as_str()))
}
GotoTarget::TypeParamTypeVarName(type_var) => {
Some(Cow::Borrowed(type_var.name.as_str()))
}
GotoTarget::TypeParamParamSpecName(spec) => Some(Cow::Borrowed(spec.name.as_str())),
GotoTarget::TypeParamTypeVarTupleName(tuple) => {
Some(Cow::Borrowed(tuple.name.as_str()))
}
GotoTarget::NonLocal { identifier, .. } => Some(Cow::Borrowed(identifier.as_str())),
GotoTarget::Globals { identifier, .. } => Some(Cow::Borrowed(identifier.as_str())),
}
}
/// Creates a `GotoTarget` from a `CoveringNode` and an offset within the node
pub(crate) fn from_covering_node<'a>(
covering_node: &crate::find_node::CoveringNode<'a>,
offset: TextSize,
) -> Option<GotoTarget<'a>> {
tracing::trace!("Covering node is of kind {:?}", covering_node.node().kind());
match covering_node.node() {
AnyNodeRef::Identifier(identifier) => match covering_node.parent() {
Some(AnyNodeRef::StmtFunctionDef(function)) => {
Some(GotoTarget::FunctionDef(function))
}
Some(AnyNodeRef::StmtClassDef(class)) => Some(GotoTarget::ClassDef(class)),
Some(AnyNodeRef::Parameter(parameter)) => Some(GotoTarget::Parameter(parameter)),
Some(AnyNodeRef::Alias(alias)) => {
// Find the containing import statement to determine the type
let import_stmt = covering_node.ancestors().find(|node| {
matches!(
node,
AnyNodeRef::StmtImport(_) | AnyNodeRef::StmtImportFrom(_)
)
});
match import_stmt {
Some(AnyNodeRef::StmtImport(_)) => {
// Regular import statement like "import x.y as z"
// Is the offset within the alias name (asname) part?
if let Some(asname) = &alias.asname {
if asname.range.contains_inclusive(offset) {
return Some(GotoTarget::ImportModuleAlias { alias });
}
}
// Is the offset in the module name part?
if alias.name.range.contains_inclusive(offset) {
let full_name = alias.name.as_str();
if let Some((component_index, component_range)) =
find_module_component(
full_name,
alias.name.range.start(),
offset,
)
{
return Some(GotoTarget::ImportModuleComponent {
module_name: full_name.to_string(),
component_index,
component_range,
});
}
}
None
}
Some(AnyNodeRef::StmtImportFrom(import_from)) => {
// From import statement like "from x import y as z"
// Is the offset within the alias name (asname) part?
if let Some(asname) = &alias.asname {
if asname.range.contains_inclusive(offset) {
return Some(GotoTarget::ImportSymbolAlias {
alias,
range: asname.range,
import_from,
});
}
}
// Is the offset in the original name part?
if alias.name.range.contains_inclusive(offset) {
return Some(GotoTarget::ImportSymbolAlias {
alias,
range: alias.name.range,
import_from,
});
}
None
}
_ => None,
}
}
Some(AnyNodeRef::StmtImportFrom(from)) => {
// Handle offset within module name in from import statements
if let Some(module_expr) = &from.module {
let full_module_name = module_expr.to_string();
if let Some((component_index, component_range)) = find_module_component(
&full_module_name,
module_expr.range.start(),
offset,
) {
return Some(GotoTarget::ImportModuleComponent {
module_name: full_module_name,
component_index,
component_range,
});
}
}
None
}
Some(AnyNodeRef::ExceptHandlerExceptHandler(handler)) => {
Some(GotoTarget::ExceptVariable(handler))
}
Some(AnyNodeRef::Keyword(keyword)) => {
// Find the containing call expression from the ancestor chain
let call_expression = covering_node
.ancestors()
.find_map(ruff_python_ast::AnyNodeRef::expr_call)?;
Some(GotoTarget::KeywordArgument {
keyword,
call_expression,
})
}
Some(AnyNodeRef::PatternMatchMapping(mapping)) => {
Some(GotoTarget::PatternMatchRest(mapping))
}
Some(AnyNodeRef::PatternKeyword(keyword)) => {
Some(GotoTarget::PatternKeywordArgument(keyword))
}
Some(AnyNodeRef::PatternMatchStar(star)) => {
Some(GotoTarget::PatternMatchStarName(star))
}
Some(AnyNodeRef::PatternMatchAs(as_pattern)) => {
Some(GotoTarget::PatternMatchAsName(as_pattern))
}
Some(AnyNodeRef::TypeParamTypeVar(var)) => {
Some(GotoTarget::TypeParamTypeVarName(var))
}
Some(AnyNodeRef::TypeParamParamSpec(bound)) => {
Some(GotoTarget::TypeParamParamSpecName(bound))
}
Some(AnyNodeRef::TypeParamTypeVarTuple(var_tuple)) => {
Some(GotoTarget::TypeParamTypeVarTupleName(var_tuple))
}
Some(AnyNodeRef::ExprAttribute(attribute)) => {
Some(GotoTarget::Expression(attribute.into()))
}
Some(AnyNodeRef::StmtNonlocal(_)) => Some(GotoTarget::NonLocal { identifier }),
Some(AnyNodeRef::StmtGlobal(_)) => Some(GotoTarget::Globals { identifier }),
None => None,
Some(parent) => {
tracing::debug!(
"Missing `GoToTarget` for identifier with parent {:?}",
parent.kind()
);
None
}
},
node => node.as_expr_ref().map(GotoTarget::Expression),
}
}
}
impl Ranged for GotoTarget<'_> {
@ -328,11 +590,7 @@ fn convert_resolved_definitions_to_targets(
}
ty_python_semantic::ResolvedDefinition::FileWithRange(file_range) => {
// For file ranges, navigate to the specific range within the file
crate::NavigationTarget {
file: file_range.file(),
focus_range: file_range.range(),
full_range: file_range.range(),
}
crate::NavigationTarget::new(file_range.file(), file_range.range())
}
})
.collect()
@ -375,145 +633,7 @@ pub(crate) fn find_goto_target(
.find_first(|node| node.is_identifier() || node.is_expression())
.ok()?;
tracing::trace!("Covering node is of kind {:?}", covering_node.node().kind());
match covering_node.node() {
AnyNodeRef::Identifier(identifier) => match covering_node.parent() {
Some(AnyNodeRef::StmtFunctionDef(function)) => Some(GotoTarget::FunctionDef(function)),
Some(AnyNodeRef::StmtClassDef(class)) => Some(GotoTarget::ClassDef(class)),
Some(AnyNodeRef::Parameter(parameter)) => Some(GotoTarget::Parameter(parameter)),
Some(AnyNodeRef::Alias(alias)) => {
// Find the containing import statement to determine the type
let import_stmt = covering_node.ancestors().find(|node| {
matches!(
node,
AnyNodeRef::StmtImport(_) | AnyNodeRef::StmtImportFrom(_)
)
});
match import_stmt {
Some(AnyNodeRef::StmtImport(_)) => {
// Regular import statement like "import x.y as z"
// Is the offset within the alias name (asname) part?
if let Some(asname) = &alias.asname {
if asname.range.contains_inclusive(offset) {
return Some(GotoTarget::ImportModuleAlias { alias });
}
}
// Is the offset in the module name part?
if alias.name.range.contains_inclusive(offset) {
let full_name = alias.name.as_str();
if let Some((component_index, component_range)) =
find_module_component(full_name, alias.name.range.start(), offset)
{
return Some(GotoTarget::ImportModuleComponent {
module_name: full_name.to_string(),
component_index,
component_range,
});
}
}
None
}
Some(AnyNodeRef::StmtImportFrom(import_from)) => {
// From import statement like "from x import y as z"
// Is the offset within the alias name (asname) part?
if let Some(asname) = &alias.asname {
if asname.range.contains_inclusive(offset) {
return Some(GotoTarget::ImportSymbolAlias {
alias,
range: asname.range,
import_from,
});
}
}
// Is the offset in the original name part?
if alias.name.range.contains_inclusive(offset) {
return Some(GotoTarget::ImportSymbolAlias {
alias,
range: alias.name.range,
import_from,
});
}
None
}
_ => None,
}
}
Some(AnyNodeRef::StmtImportFrom(from)) => {
// Handle offset within module name in from import statements
if let Some(module_expr) = &from.module {
let full_module_name = module_expr.to_string();
if let Some((component_index, component_range)) =
find_module_component(&full_module_name, module_expr.range.start(), offset)
{
return Some(GotoTarget::ImportModuleComponent {
module_name: full_module_name,
component_index,
component_range,
});
}
}
None
}
Some(AnyNodeRef::ExceptHandlerExceptHandler(handler)) => {
Some(GotoTarget::ExceptVariable(handler))
}
Some(AnyNodeRef::Keyword(keyword)) => {
// Find the containing call expression from the ancestor chain
let call_expression = covering_node
.ancestors()
.find_map(ruff_python_ast::AnyNodeRef::expr_call)?;
Some(GotoTarget::KeywordArgument {
keyword,
call_expression,
})
}
Some(AnyNodeRef::PatternMatchMapping(mapping)) => {
Some(GotoTarget::PatternMatchRest(mapping))
}
Some(AnyNodeRef::PatternKeyword(keyword)) => {
Some(GotoTarget::PatternKeywordArgument(keyword))
}
Some(AnyNodeRef::PatternMatchStar(star)) => {
Some(GotoTarget::PatternMatchStarName(star))
}
Some(AnyNodeRef::PatternMatchAs(as_pattern)) => {
Some(GotoTarget::PatternMatchAsName(as_pattern))
}
Some(AnyNodeRef::TypeParamTypeVar(var)) => Some(GotoTarget::TypeParamTypeVarName(var)),
Some(AnyNodeRef::TypeParamParamSpec(bound)) => {
Some(GotoTarget::TypeParamParamSpecName(bound))
}
Some(AnyNodeRef::TypeParamTypeVarTuple(var_tuple)) => {
Some(GotoTarget::TypeParamTypeVarTupleName(var_tuple))
}
Some(AnyNodeRef::ExprAttribute(attribute)) => {
Some(GotoTarget::Expression(attribute.into()))
}
Some(AnyNodeRef::StmtNonlocal(_)) => Some(GotoTarget::NonLocal { identifier }),
Some(AnyNodeRef::StmtGlobal(_)) => Some(GotoTarget::Globals { identifier }),
None => None,
Some(parent) => {
tracing::debug!(
"Missing `GoToTarget` for identifier with parent {:?}",
parent.kind()
);
None
}
},
node => node.as_expr_ref().map(GotoTarget::Expression),
}
GotoTarget::from_covering_node(&covering_node, offset)
}
/// Helper function to resolve a module name and create a navigation target.
@ -526,11 +646,9 @@ fn resolve_module_to_navigation_target(
if let Some(module_name) = ModuleName::new(module_name_str) {
if let Some(resolved_module) = resolve_module(db, &module_name) {
if let Some(module_file) = resolved_module.file(db) {
return Some(crate::NavigationTargets::single(crate::NavigationTarget {
file: module_file,
focus_range: TextRange::default(),
full_range: TextRange::default(),
}));
return Some(crate::NavigationTargets::single(
crate::NavigationTarget::new(module_file, TextRange::default()),
));
}
}
}

View file

@ -8,6 +8,7 @@ mod goto_type_definition;
mod hover;
mod inlay_hints;
mod markup;
mod references;
mod semantic_tokens;
mod signature_help;
mod stub_mapping;
@ -18,6 +19,7 @@ pub use goto::{goto_declaration, goto_definition, goto_type_definition};
pub use hover::hover;
pub use inlay_hints::inlay_hints;
pub use markup::MarkupKind;
pub use references::references;
pub use semantic_tokens::{
SemanticToken, SemanticTokenModifier, SemanticTokenType, SemanticTokens, semantic_tokens,
};
@ -86,6 +88,15 @@ pub struct NavigationTarget {
}
impl NavigationTarget {
/// Creates a new `NavigationTarget` where the focus and full range are identical.
pub fn new(file: File, range: TextRange) -> Self {
Self {
file,
focus_range: range,
full_range: range,
}
}
pub fn file(&self) -> File {
self.file
}
@ -291,6 +302,7 @@ mod tests {
));
let mut cursor: Option<Cursor> = None;
for &Source {
ref path,
ref contents,
@ -299,11 +311,10 @@ mod tests {
{
db.write_file(path, contents)
.expect("write to memory file system to be successful");
let Some(offset) = cursor_offset else {
continue;
};
let file = system_path_to_file(&db, path).expect("newly written file to existing");
if let Some(offset) = cursor_offset {
// This assert should generally never trip, since
// we have an assert on `CursorTestBuilder::source`
// to ensure we never have more than one marker.
@ -313,6 +324,7 @@ mod tests {
);
cursor = Some(Cursor { file, offset });
}
}
let search_paths = SearchPathSettings::new(vec![SystemPathBuf::from("/")])
.to_search_paths(db.system(), db.vendored())

File diff suppressed because it is too large Load diff

View file

@ -209,6 +209,7 @@ impl Server {
type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
definition_provider: Some(lsp_types::OneOf::Left(true)),
declaration_provider: Some(DeclarationCapability::Simple(true)),
references_provider: Some(lsp_types::OneOf::Left(true)),
hover_provider: Some(HoverProviderCapability::Simple(true)),
signature_help_provider: Some(SignatureHelpOptions {
trigger_characters: Some(vec!["(".to_string(), ",".to_string()]),

View file

@ -56,6 +56,9 @@ pub(super) fn request(req: server::Request) -> Task {
requests::HoverRequestHandler::METHOD => background_document_request_task::<
requests::HoverRequestHandler,
>(req, BackgroundSchedule::Worker),
requests::ReferencesRequestHandler::METHOD => background_document_request_task::<
requests::ReferencesRequestHandler,
>(req, BackgroundSchedule::Worker),
requests::InlayHintRequestHandler::METHOD => background_document_request_task::<
requests::InlayHintRequestHandler,
>(req, BackgroundSchedule::Worker),

View file

@ -5,6 +5,7 @@ mod goto_definition;
mod goto_type_definition;
mod hover;
mod inlay_hints;
mod references;
mod semantic_tokens;
mod semantic_tokens_range;
mod shutdown;
@ -18,6 +19,7 @@ pub(super) use goto_definition::GotoDefinitionRequestHandler;
pub(super) use goto_type_definition::GotoTypeDefinitionRequestHandler;
pub(super) use hover::HoverRequestHandler;
pub(super) use inlay_hints::InlayHintRequestHandler;
pub(super) use references::ReferencesRequestHandler;
pub(super) use semantic_tokens::SemanticTokensRequestHandler;
pub(super) use semantic_tokens_range::SemanticTokensRangeRequestHandler;
pub(super) use shutdown::ShutdownHandler;

View file

@ -0,0 +1,65 @@
use std::borrow::Cow;
use lsp_types::request::References;
use lsp_types::{Location, ReferenceParams, Url};
use ruff_db::source::{line_index, source_text};
use ty_ide::references;
use ty_project::ProjectDatabase;
use crate::document::{PositionExt, ToLink};
use crate::server::api::traits::{
BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
};
use crate::session::DocumentSnapshot;
use crate::session::client::Client;
pub(crate) struct ReferencesRequestHandler;
impl RequestHandler for ReferencesRequestHandler {
type RequestType = References;
}
impl BackgroundDocumentRequestHandler for ReferencesRequestHandler {
fn document_url(params: &ReferenceParams) -> Cow<Url> {
Cow::Borrowed(&params.text_document_position.text_document.uri)
}
fn run_with_snapshot(
db: &ProjectDatabase,
snapshot: DocumentSnapshot,
_client: &Client,
params: ReferenceParams,
) -> crate::server::Result<Option<Vec<Location>>> {
if snapshot.client_settings().is_language_services_disabled() {
return Ok(None);
}
let Some(file) = snapshot.file(db) else {
return Ok(None);
};
let source = source_text(db, file);
let line_index = line_index(db, file);
let offset = params.text_document_position.position.to_text_size(
&source,
&line_index,
snapshot.encoding(),
);
let include_declaration = params.context.include_declaration;
let Some(references_result) = references(db, file, offset, include_declaration) else {
return Ok(None);
};
let locations: Vec<_> = references_result
.into_iter()
.flat_map(|ranged| ranged.value.into_iter())
.filter_map(|target| target.to_location(db, snapshot.encoding()))
.collect();
Ok(Some(locations))
}
}
impl RetriableRequestHandler for ReferencesRequestHandler {}

View file

@ -26,6 +26,7 @@ expression: initialization_result
},
"definitionProvider": true,
"typeDefinitionProvider": true,
"referencesProvider": true,
"declarationProvider": true,
"semanticTokensProvider": {
"legend": {

View file

@ -26,6 +26,7 @@ expression: initialization_result
},
"definitionProvider": true,
"typeDefinitionProvider": true,
"referencesProvider": true,
"declarationProvider": true,
"semanticTokensProvider": {
"legend": {