[ty] Improve lsp handling of hover/goto on imports (#21572)

* Fixes https://github.com/astral-sh/ty/issues/1011
* Also fixes the fact that we didn't handle `.x` properly *at all* in
hover/goto

It turns out all of our import handling completely ignored the `level`
(number of relative `.`'s) in a `from ..x.y import z` statement. It was
nice seeing how much my understanding of `ty` has improved -- previously
this would have all been opaque to me but now it was just, completely
glaring and blatant.

Fixing this required refactoring all the import code to take the
importing file into consideration. I ended up refactoring a bunch of
code to pass around/require `SemanticModel` more, as it's the natural
API for resolving this kind of import (it actually had an API for this
that was just... dead code, whoops!).
This commit is contained in:
Aria Desires 2025-11-22 11:06:16 -05:00 committed by GitHub
parent 3410041b4c
commit 859f9ec21a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 403 additions and 89 deletions

View File

@ -61,6 +61,7 @@ pub(crate) enum GotoTarget<'a> {
/// ```
ImportModuleComponent {
module_name: String,
level: u32,
component_index: usize,
component_range: TextRange,
},
@ -302,12 +303,21 @@ impl GotoTarget<'_> {
// (i.e. the type of `MyClass` in `MyClass()` is `<class MyClass>` and not `() -> MyClass`)
GotoTarget::Call { callable, .. } => callable.inferred_type(model),
GotoTarget::TypeParamTypeVarName(typevar) => typevar.inferred_type(model),
GotoTarget::ImportModuleComponent {
module_name,
component_index,
level,
..
} => {
// We don't currently support hovering the bare `.` so there is always a name
let module = import_name(module_name, *component_index);
model.resolve_module_type(Some(module), *level)?
}
// TODO: Support identifier targets
GotoTarget::PatternMatchRest(_)
| GotoTarget::PatternKeywordArgument(_)
| GotoTarget::PatternMatchStarName(_)
| GotoTarget::PatternMatchAsName(_)
| GotoTarget::ImportModuleComponent { .. }
| GotoTarget::TypeParamParamSpecName(_)
| GotoTarget::TypeParamTypeVarTupleName(_)
| GotoTarget::NonLocal { .. }
@ -353,37 +363,30 @@ impl GotoTarget<'_> {
/// as just returning a raw `NavigationTarget`.
pub(crate) fn get_definition_targets<'db>(
&self,
file: ruff_db::files::File,
db: &'db dyn crate::Db,
model: &SemanticModel<'db>,
alias_resolution: ImportAliasResolution,
) -> Option<DefinitionsOrTargets<'db>> {
use crate::NavigationTarget;
let db = model.db();
let file = model.file();
match self {
GotoTarget::Expression(expression) => definitions_for_expression(db, file, expression)
.map(DefinitionsOrTargets::Definitions),
GotoTarget::Expression(expression) => {
definitions_for_expression(model, expression).map(DefinitionsOrTargets::Definitions)
}
// For already-defined symbols, they are their own definitions
GotoTarget::FunctionDef(function) => {
let model = SemanticModel::new(db, file);
Some(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Definition(function.definition(&model)),
]))
}
GotoTarget::FunctionDef(function) => Some(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Definition(function.definition(model)),
])),
GotoTarget::ClassDef(class) => {
let model = SemanticModel::new(db, file);
Some(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Definition(class.definition(&model)),
]))
}
GotoTarget::ClassDef(class) => Some(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Definition(class.definition(model)),
])),
GotoTarget::Parameter(parameter) => {
let model = SemanticModel::new(db, file);
Some(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Definition(parameter.definition(&model)),
]))
}
GotoTarget::Parameter(parameter) => Some(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Definition(parameter.definition(model)),
])),
// For import aliases (offset within 'y' or 'z' in "from x import y as z")
GotoTarget::ImportSymbolAlias {
@ -404,24 +407,18 @@ impl GotoTarget<'_> {
GotoTarget::ImportModuleComponent {
module_name,
component_index,
level,
..
} => {
// Handle both `import foo.bar` and `from foo.bar import baz` where offset is within module component
let components: Vec<&str> = module_name.split('.').collect();
// Build the module name up to and including the component containing the offset
let target_module_name = components[..=*component_index].join(".");
// Try to resolve the module
definitions_for_module(db, &target_module_name)
// We don't currently support hovering the bare `.` so there is always a name
let module = import_name(module_name, *component_index);
definitions_for_module(model, Some(module), *level)
}
// Handle import aliases (offset within 'z' in "import x.y as z")
GotoTarget::ImportModuleAlias { alias } => {
if alias_resolution == ImportAliasResolution::ResolveAliases {
let full_module_name = alias.name.as_str();
// Try to resolve the module
definitions_for_module(db, full_module_name)
definitions_for_module(model, Some(alias.name.as_str()), 0)
} else {
let alias_range = alias.asname.as_ref().unwrap().range;
Some(DefinitionsOrTargets::Targets(
@ -444,9 +441,8 @@ impl GotoTarget<'_> {
// For exception variables, they are their own definitions (like parameters)
GotoTarget::ExceptVariable(except_handler) => {
let model = SemanticModel::new(db, file);
Some(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Definition(except_handler.definition(&model)),
ResolvedDefinition::Definition(except_handler.definition(model)),
]))
}
@ -478,9 +474,9 @@ impl GotoTarget<'_> {
//
// Prefer the function impl over the callable so that its docstrings win if defined.
GotoTarget::Call { callable, call } => {
let mut definitions = definitions_for_callable(db, file, call);
let mut definitions = definitions_for_callable(model, call);
let expr_definitions =
definitions_for_expression(db, file, callable).unwrap_or_default();
definitions_for_expression(model, callable).unwrap_or_default();
definitions.extend(expr_definitions);
if definitions.is_empty() {
@ -491,18 +487,15 @@ impl GotoTarget<'_> {
}
GotoTarget::BinOp { expression, .. } => {
let model = SemanticModel::new(db, file);
let (definitions, _) =
ty_python_semantic::definitions_for_bin_op(db, &model, expression)?;
ty_python_semantic::definitions_for_bin_op(db, model, expression)?;
Some(DefinitionsOrTargets::Definitions(definitions))
}
GotoTarget::UnaryOp { expression, .. } => {
let model = SemanticModel::new(db, file);
let (definitions, _) =
ty_python_semantic::definitions_for_unary_op(db, &model, expression)?;
ty_python_semantic::definitions_for_unary_op(db, model, expression)?;
Some(DefinitionsOrTargets::Definitions(definitions))
}
@ -632,6 +625,7 @@ impl GotoTarget<'_> {
{
return Some(GotoTarget::ImportModuleComponent {
module_name: full_name.to_string(),
level: 0,
component_index,
component_range,
});
@ -672,14 +666,12 @@ impl GotoTarget<'_> {
// 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,
) {
if let Some((component_index, component_range)) =
find_module_component(&full_module_name, module_expr.start(), offset)
{
return Some(GotoTarget::ImportModuleComponent {
module_name: full_module_name,
level: from.level,
component_index,
component_range,
});
@ -876,27 +868,26 @@ fn convert_resolved_definitions_to_targets(
/// Shared helper to get definitions for an expr (that is presumably a name/attr)
fn definitions_for_expression<'db>(
db: &'db dyn crate::Db,
file: ruff_db::files::File,
model: &SemanticModel<'db>,
expression: &ruff_python_ast::ExprRef<'_>,
) -> Option<Vec<ResolvedDefinition<'db>>> {
match expression {
ast::ExprRef::Name(name) => Some(definitions_for_name(db, file, name)),
ast::ExprRef::Name(name) => Some(definitions_for_name(model.db(), model.file(), name)),
ast::ExprRef::Attribute(attribute) => Some(ty_python_semantic::definitions_for_attribute(
db, file, attribute,
model.db(),
model.file(),
attribute,
)),
_ => None,
}
}
fn definitions_for_callable<'db>(
db: &'db dyn crate::Db,
file: ruff_db::files::File,
model: &SemanticModel<'db>,
call: &ast::ExprCall,
) -> Vec<ResolvedDefinition<'db>> {
let model = SemanticModel::new(db, file);
// Attempt to refine to a specific call
let signature_info = call_signature_details(db, &model, call);
let signature_info = call_signature_details(model.db(), model, call);
signature_info
.into_iter()
.filter_map(|signature| signature.definition.map(ResolvedDefinition::Definition))
@ -947,7 +938,9 @@ pub(crate) fn find_goto_target(
}
let covering_node = covering_node(parsed.syntax().into(), token.range())
.find_first(|node| node.is_identifier() || node.is_expression())
.find_first(|node| {
node.is_identifier() || node.is_expression() || node.is_stmt_import_from()
})
.ok()?;
GotoTarget::from_covering_node(&covering_node, offset, parsed.tokens())
@ -955,21 +948,15 @@ pub(crate) fn find_goto_target(
/// Helper function to resolve a module name and create a navigation target.
fn definitions_for_module<'db>(
db: &'db dyn crate::Db,
module_name_str: &str,
model: &SemanticModel,
module: Option<&str>,
level: u32,
) -> Option<DefinitionsOrTargets<'db>> {
use ty_python_semantic::{ModuleName, resolve_module};
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(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Module(module_file),
]));
}
}
}
None
let module = model.resolve_module(module, level)?;
let file = module.file(model.db())?;
Some(DefinitionsOrTargets::Definitions(vec![
ResolvedDefinition::Module(file),
]))
}
/// Helper function to extract module component information from a dotted module name
@ -983,9 +970,7 @@ fn find_module_component(
// Split the module name into components and find which one contains the offset
let mut current_pos = 0;
let components: Vec<&str> = full_module_name.split('.').collect();
for (i, component) in components.iter().enumerate() {
for (i, component) in full_module_name.split('.').enumerate() {
let component_start = current_pos;
let component_end = current_pos + component.len();
@ -1004,3 +989,16 @@ fn find_module_component(
None
}
/// Helper to get the module name up to the given component index
fn import_name(module_name: &str, component_index: usize) -> &str {
// We want everything to the left of the nth `.`
// If there's no nth `.` then we want the whole thing.
let idx = module_name
.match_indices('.')
.nth(component_index)
.map(|(idx, _)| idx)
.unwrap_or(module_name.len());
&module_name[..idx]
}

View File

@ -3,7 +3,7 @@ use crate::{Db, NavigationTargets, RangedValue};
use ruff_db::files::{File, FileRange};
use ruff_db::parsed::parsed_module;
use ruff_text_size::{Ranged, TextSize};
use ty_python_semantic::ImportAliasResolution;
use ty_python_semantic::{ImportAliasResolution, SemanticModel};
/// Navigate to the declaration of a symbol.
///
@ -18,8 +18,9 @@ pub fn goto_declaration(
let module = parsed_module(db, file).load(db);
let goto_target = find_goto_target(&module, offset)?;
let model = SemanticModel::new(db, file);
let declaration_targets = goto_target
.get_definition_targets(file, db, ImportAliasResolution::ResolveAliases)?
.get_definition_targets(&model, ImportAliasResolution::ResolveAliases)?
.declaration_targets(db)?;
Some(RangedValue {

View File

@ -3,7 +3,7 @@ use crate::{Db, NavigationTargets, RangedValue};
use ruff_db::files::{File, FileRange};
use ruff_db::parsed::parsed_module;
use ruff_text_size::{Ranged, TextSize};
use ty_python_semantic::ImportAliasResolution;
use ty_python_semantic::{ImportAliasResolution, SemanticModel};
/// Navigate to the definition of a symbol.
///
@ -18,9 +18,9 @@ pub fn goto_definition(
) -> Option<RangedValue<NavigationTargets>> {
let module = parsed_module(db, file).load(db);
let goto_target = find_goto_target(&module, offset)?;
let model = SemanticModel::new(db, file);
let definition_targets = goto_target
.get_definition_targets(file, db, ImportAliasResolution::ResolveAliases)?
.get_definition_targets(&model, ImportAliasResolution::ResolveAliases)?
.definition_targets(db)?;
Some(RangedValue {

View File

@ -285,6 +285,300 @@ mod tests {
");
}
#[test]
fn goto_type_of_import_module() {
let mut test = cursor_test(
r#"
import l<CURSOR>ib
"#,
);
test.write_file("lib.py", "a = 10").unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib.py:1:1
|
1 | a = 10
| ^^^^^^
|
info: Source
--> main.py:2:8
|
2 | import lib
| ^^^
|
");
}
#[test]
fn goto_type_of_import_module_multi1() {
let mut test = cursor_test(
r#"
import li<CURSOR>b.submod
"#,
);
test.write_file("lib/__init__.py", "b = 7").unwrap();
test.write_file("lib/submod.py", "a = 10").unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib/__init__.py:1:1
|
1 | b = 7
| ^^^^^
|
info: Source
--> main.py:2:8
|
2 | import lib.submod
| ^^^
|
");
}
#[test]
fn goto_type_of_import_module_multi2() {
let mut test = cursor_test(
r#"
import lib.subm<CURSOR>od
"#,
);
test.write_file("lib/__init__.py", "b = 7").unwrap();
test.write_file("lib/submod.py", "a = 10").unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib/submod.py:1:1
|
1 | a = 10
| ^^^^^^
|
info: Source
--> main.py:2:12
|
2 | import lib.submod
| ^^^^^^
|
");
}
#[test]
fn goto_type_of_from_import_module() {
let mut test = cursor_test(
r#"
from l<CURSOR>ib import a
"#,
);
test.write_file("lib.py", "a = 10").unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib.py:1:1
|
1 | a = 10
| ^^^^^^
|
info: Source
--> main.py:2:6
|
2 | from lib import a
| ^^^
|
");
}
#[test]
fn goto_type_of_from_import_module_multi1() {
let mut test = cursor_test(
r#"
from li<CURSOR>b.submod import a
"#,
);
test.write_file("lib/__init__.py", "b = 7").unwrap();
test.write_file("lib/submod.py", "a = 10").unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib/__init__.py:1:1
|
1 | b = 7
| ^^^^^
|
info: Source
--> main.py:2:6
|
2 | from lib.submod import a
| ^^^
|
");
}
#[test]
fn goto_type_of_from_import_module_multi2() {
let mut test = cursor_test(
r#"
from lib.subm<CURSOR>od import a
"#,
);
test.write_file("lib/__init__.py", "b = 7").unwrap();
test.write_file("lib/submod.py", "a = 10").unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib/submod.py:1:1
|
1 | a = 10
| ^^^^^^
|
info: Source
--> main.py:2:10
|
2 | from lib.submod import a
| ^^^^^^
|
");
}
#[test]
fn goto_type_of_from_import_rel1() {
let mut test = CursorTest::builder()
.source(
"lib/sub/__init__.py",
r#"
from .bot.bot<CURSOR>mod import *
sub = 2
"#,
)
.build();
test.write_file("lib/__init__.py", "lib = 1").unwrap();
// test.write_file("lib/sub/__init__.py", "sub = 2").unwrap();
test.write_file("lib/sub/bot/__init__.py", "bot = 3")
.unwrap();
test.write_file("lib/sub/submod.py", "submod = 21").unwrap();
test.write_file("lib/sub/bot/botmod.py", "botmod = 31")
.unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib/sub/bot/botmod.py:1:1
|
1 | botmod = 31
| ^^^^^^^^^^^
|
info: Source
--> lib/sub/__init__.py:2:11
|
2 | from .bot.botmod import *
| ^^^^^^
3 | sub = 2
|
");
}
#[test]
fn goto_type_of_from_import_rel2() {
let mut test = CursorTest::builder()
.source(
"lib/sub/__init__.py",
r#"
from .bo<CURSOR>t.botmod import *
sub = 2
"#,
)
.build();
test.write_file("lib/__init__.py", "lib = 1").unwrap();
// test.write_file("lib/sub/__init__.py", "sub = 2").unwrap();
test.write_file("lib/sub/bot/__init__.py", "bot = 3")
.unwrap();
test.write_file("lib/sub/submod.py", "submod = 21").unwrap();
test.write_file("lib/sub/bot/botmod.py", "botmod = 31")
.unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib/sub/bot/__init__.py:1:1
|
1 | bot = 3
| ^^^^^^^
|
info: Source
--> lib/sub/__init__.py:2:7
|
2 | from .bot.botmod import *
| ^^^
3 | sub = 2
|
");
}
#[test]
fn goto_type_of_from_import_rel3() {
let mut test = CursorTest::builder()
.source(
"lib/sub/__init__.py",
r#"
from .<CURSOR>bot.botmod import *
sub = 2
"#,
)
.build();
test.write_file("lib/__init__.py", "lib = 1").unwrap();
// test.write_file("lib/sub/__init__.py", "sub = 2").unwrap();
test.write_file("lib/sub/bot/__init__.py", "bot = 3")
.unwrap();
test.write_file("lib/sub/submod.py", "submod = 21").unwrap();
test.write_file("lib/sub/bot/botmod.py", "botmod = 31")
.unwrap();
assert_snapshot!(test.goto_type_definition(), @r"
info[goto-type-definition]: Type definition
--> lib/sub/bot/__init__.py:1:1
|
1 | bot = 3
| ^^^^^^^
|
info: Source
--> lib/sub/__init__.py:2:7
|
2 | from .bot.botmod import *
| ^^^
3 | sub = 2
|
");
}
#[test]
fn goto_type_of_from_import_rel4() {
let mut test = CursorTest::builder()
.source(
"lib/sub/__init__.py",
r#"
from .<CURSOR> import submod
sub = 2
"#,
)
.build();
test.write_file("lib/__init__.py", "lib = 1").unwrap();
// test.write_file("lib/sub/__init__.py", "sub = 2").unwrap();
test.write_file("lib/sub/bot/__init__.py", "bot = 3")
.unwrap();
test.write_file("lib/sub/submod.py", "submod = 21").unwrap();
test.write_file("lib/sub/bot/botmod.py", "botmod = 31")
.unwrap();
assert_snapshot!(test.goto_type_definition(), @"No goto target found");
}
#[test]
fn goto_type_of_expression_with_module() {
let mut test = cursor_test(

View File

@ -22,8 +22,7 @@ pub fn hover(db: &dyn Db, file: File, offset: TextSize) -> Option<RangedValue<Ho
let model = SemanticModel::new(db, file);
let docs = goto_target
.get_definition_targets(
file,
db,
&model,
ty_python_semantic::ImportAliasResolution::ResolveAliases,
)
.and_then(|definitions| definitions.docstring(db))
@ -1488,11 +1487,17 @@ def ab(a: int, *, c: int):
.unwrap();
assert_snapshot!(test.hover(), @r"
<module 'lib'>
---------------------------------------------
The cool lib_py module!
Wow this module rocks.
---------------------------------------------
```python
<module 'lib'>
```
---
The cool lib/_py module!
Wow this module rocks.

View File

@ -20,7 +20,7 @@ use ruff_python_ast::{
};
use ruff_python_parser::Tokens;
use ruff_text_size::{Ranged, TextRange};
use ty_python_semantic::ImportAliasResolution;
use ty_python_semantic::{ImportAliasResolution, SemanticModel};
/// Mode for references search behavior
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@ -48,8 +48,9 @@ pub(crate) fn references(
// Get the definitions for the symbol at the cursor position
// When finding references, do not resolve any local aliases.
let model = SemanticModel::new(db, file);
let target_definitions_nav = goto_target
.get_definition_targets(file, db, ImportAliasResolution::PreserveAliases)?
.get_definition_targets(&model, ImportAliasResolution::PreserveAliases)?
.definition_targets(db)?;
let target_definitions: Vec<NavigationTarget> = target_definitions_nav.into_iter().collect();
@ -289,8 +290,9 @@ impl LocalReferencesFinder<'_> {
GotoTarget::from_covering_node(covering_node, offset, self.tokens)
{
// Get the definitions for this goto target
let model = SemanticModel::new(self.db, self.file);
if let Some(current_definitions_nav) = goto_target
.get_definition_targets(self.file, self.db, ImportAliasResolution::PreserveAliases)
.get_definition_targets(&model, ImportAliasResolution::PreserveAliases)
.and_then(|definitions| definitions.declaration_targets(self.db))
{
let current_definitions: Vec<NavigationTarget> =

View File

@ -3,12 +3,13 @@ use crate::references::{ReferencesMode, references};
use crate::{Db, ReferenceTarget};
use ruff_db::files::File;
use ruff_text_size::{Ranged, TextSize};
use ty_python_semantic::ImportAliasResolution;
use ty_python_semantic::{ImportAliasResolution, SemanticModel};
/// Returns the range of the symbol if it can be renamed, None if not.
pub fn can_rename(db: &dyn Db, file: File, offset: TextSize) -> Option<ruff_text_size::TextRange> {
let parsed = ruff_db::parsed::parsed_module(db, file);
let module = parsed.load(db);
let model = SemanticModel::new(db, file);
// Get the definitions for the symbol at the offset
let goto_target = find_goto_target(&module, offset)?;
@ -24,7 +25,7 @@ pub fn can_rename(db: &dyn Db, file: File, offset: TextSize) -> Option<ruff_text
let current_file_in_project = is_file_in_project(db, file);
if let Some(definition_targets) = goto_target
.get_definition_targets(file, db, ImportAliasResolution::PreserveAliases)
.get_definition_targets(&model, ImportAliasResolution::PreserveAliases)
.and_then(|definitions| definitions.declaration_targets(db))
{
for target in &definition_targets {

View File

@ -31,6 +31,10 @@ impl<'db> SemanticModel<'db> {
self.db
}
pub fn file(&self) -> File {
self.file
}
pub fn file_path(&self) -> &FilePath {
self.file.path(self.db)
}
@ -70,8 +74,17 @@ impl<'db> SemanticModel<'db> {
members
}
pub fn resolve_module(&self, module_name: &ModuleName) -> Option<Module<'_>> {
resolve_module(self.db, module_name)
/// Resolve the given import made in this file to a Type
pub fn resolve_module_type(&self, module: Option<&str>, level: u32) -> Option<Type<'db>> {
let module = self.resolve_module(module, level)?;
Some(Type::module_literal(self.db, self.file, module))
}
/// Resolve the given import made in this file to a Module
pub fn resolve_module(&self, module: Option<&str>, level: u32) -> Option<Module<'db>> {
let module_name =
ModuleName::from_identifier_parts(self.db, self.file, module, level).ok()?;
resolve_module(self.db, &module_name)
}
/// Returns completions for symbols available in a `import <CURSOR>` context.