fix spans of ImportFromSubmodule

This commit is contained in:
Aria Desires 2025-12-04 12:03:08 -05:00
parent cccb0bbaa4
commit a3744087c6
3 changed files with 63 additions and 19 deletions

View File

@ -2622,16 +2622,12 @@ def ab(a: int, *, c: int): ...
)
.build();
// TODO(submodule-imports): this should only highlight `subpkg` in the import statement
// This happens because DefinitionKind::ImportFromSubmodule claims the entire ImportFrom node,
// which is correct but unhelpful. Unfortunately even if it only claimed the LHS identifier it
// would highlight `subpkg.submod` which is strictly better but still isn't what we want.
assert_snapshot!(test.goto_declaration(), @r"
info[goto-declaration]: Declaration
--> mypackage/__init__.py:2:1
--> mypackage/__init__.py:2:7
|
2 | from .subpkg.submod import val
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
| ^^^^^^
3 |
4 | x = subpkg
|
@ -2846,25 +2842,26 @@ def ab(a: int, *, c: int): ...
)
.build();
// TODO(submodule-imports): Ok this one is FASCINATING and it's kinda right but confusing!
// TODO(submodule-imports): Ok this one is FASCINATING but definitely wrong!
//
// So there's 3 relevant definitions here:
//
// * `subpkg: int = 10` in the other file is in fact the original definition
// * `subpkg: int = 10` in the other file is in fact the original definition.
// Including it here is accurate and possibly useful?
//
// * the LHS `subpkg` in the import is an instance of `subpkg = ...`
// because it's a `DefinitionKind::ImportFromSubmodle`.
// This is the span that covers the entire import.
// Including it here is Pedantically Correct but Unhelpful.
//
// * `the RHS `subpkg` in the import is a second instance of `subpkg = ...`
// that *immediately* overwrites the `ImportFromSubmodule`'s definition
// This span seemingly doesn't appear at all!? Is it getting hidden by the LHS span?
// that *immediately* overwrites the `ImportFromSubmodule`'s definition.
// This is the most important one and doesn't show up at all! Sadness!
assert_snapshot!(test.goto_declaration(), @r"
info[goto-declaration]: Declaration
--> mypackage/__init__.py:2:1
--> mypackage/__init__.py:2:7
|
2 | from .subpkg import subpkg
| ^^^^^^^^^^^^^^^^^^^^^^^^^^
| ^^^^^^
3 |
4 | x = subpkg
|

View File

@ -1529,7 +1529,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
// Record whether this is equivalent to `from . import ...`
is_self_import = module_name == thispackage;
if node.module.is_some()
if let Some(module_node) = &node.module
&& let Some(relative_submodule) = module_name.relative_to(&thispackage)
&& let Some(direct_submodule) = relative_submodule.components().next()
&& !self.seen_submodule_imports.contains(direct_submodule)
@ -1540,9 +1540,23 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
let direct_submodule_name = Name::new(direct_submodule);
let symbol = self.add_symbol(direct_submodule_name);
let module_index = if node.level == 0 {
// "whatever.thispackage.x.y" we want `x`
thispackage.components().count()
} else {
// ".x.y" we want `x` (level 1 => index 0)
// "..x.y" we want `y` (level 2 => index 1)
// (The Identifier doesn't include the prefix dots)
node.level as usize - 1
};
self.add_definition(
symbol.into(),
ImportFromSubmoduleDefinitionNodeRef { node },
ImportFromSubmoduleDefinitionNodeRef {
node,
module: module_node,
module_index,
},
);
}
}

View File

@ -3,7 +3,7 @@ use std::ops::Deref;
use ruff_db::files::{File, FileRange};
use ruff_db::parsed::{ParsedModuleRef, parsed_module};
use ruff_python_ast as ast;
use ruff_text_size::{Ranged, TextRange};
use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
use crate::Db;
use crate::ast_node_ref::AstNodeRef;
@ -374,6 +374,8 @@ pub(crate) struct ImportFromDefinitionNodeRef<'ast> {
#[derive(Copy, Clone, Debug)]
pub(crate) struct ImportFromSubmoduleDefinitionNodeRef<'ast> {
pub(crate) node: &'ast ast::StmtImportFrom,
pub(crate) module: &'ast ast::Identifier,
pub(crate) module_index: usize,
}
#[derive(Copy, Clone, Debug)]
@ -456,8 +458,12 @@ impl<'db> DefinitionNodeRef<'_, 'db> {
}),
DefinitionNodeRef::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef {
node,
module,
module_index,
}) => DefinitionKind::ImportFromSubmodule(ImportFromSubmoduleDefinitionKind {
node: AstNodeRef::new(parsed, node),
module: AstNodeRef::new(parsed, module),
module_index,
}),
DefinitionNodeRef::ImportStar(star_import) => {
let StarImportDefinitionNodeRef { node, symbol_id } = star_import;
@ -584,7 +590,9 @@ impl<'db> DefinitionNodeRef<'_, 'db> {
alias_index,
is_reexported: _,
}) => (&node.names[alias_index]).into(),
Self::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef { node }) => node.into(),
Self::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef { node, .. }) => {
node.into()
}
// INVARIANT: for an invalid-syntax statement such as `from foo import *, bar, *`,
// we only create a `StarImportDefinitionKind` for the *first* `*` alias in the names list.
Self::ImportStar(StarImportDefinitionNodeRef { node, symbol_id: _ }) => node
@ -755,7 +763,7 @@ impl DefinitionKind<'_> {
match self {
DefinitionKind::Import(import) => import.alias(module).range(),
DefinitionKind::ImportFrom(import) => import.alias(module).range(),
DefinitionKind::ImportFromSubmodule(import) => import.import(module).range(),
DefinitionKind::ImportFromSubmodule(import) => import.target_range(module),
DefinitionKind::StarImport(import) => import.alias(module).range(),
DefinitionKind::Function(function) => function.node(module).name.range(),
DefinitionKind::Class(class) => class.node(module).name.range(),
@ -793,7 +801,7 @@ impl DefinitionKind<'_> {
match self {
DefinitionKind::Import(import) => import.alias(module).range(),
DefinitionKind::ImportFrom(import) => import.alias(module).range(),
DefinitionKind::ImportFromSubmodule(import) => import.import(module).range(),
DefinitionKind::ImportFromSubmodule(import) => import.module(module).range(),
DefinitionKind::StarImport(import) => import.import(module).range(),
DefinitionKind::Function(function) => function.node(module).range(),
DefinitionKind::Class(class) => class.node(module).range(),
@ -1033,12 +1041,37 @@ impl ImportFromDefinitionKind {
#[derive(Clone, Debug, get_size2::GetSize)]
pub struct ImportFromSubmoduleDefinitionKind {
node: AstNodeRef<ast::StmtImportFrom>,
module: AstNodeRef<ast::Identifier>,
module_index: usize,
}
impl ImportFromSubmoduleDefinitionKind {
pub fn import<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::StmtImportFrom {
self.node.node(module)
}
pub fn module<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::Identifier {
self.module.node(module)
}
pub fn target_range(&self, module: &ParsedModuleRef) -> TextRange {
let module_ident = self.module(module);
let module_str = module_ident.as_str();
let Some(component_str) = module_str.split('.').nth(self.module_index) else {
// This shouldn't happen but just in case, provide a safe default
return module_ident.range();
};
let base_addr = module_str.as_ptr().addr();
let component_addr = component_str.as_ptr().addr();
let offset = base_addr.saturating_sub(component_addr);
let Ok(offset_size) = TextSize::try_from(offset) else {
// This shouldn't happen but just in case, provide a safe default
return module_ident.range();
};
let start = module_ident.start().saturating_add(offset_size);
let end = start.saturating_add(component_str.text_len());
TextRange::new(start, end)
}
}
#[derive(Clone, Debug, get_size2::GetSize)]