From 1af318534a01ba6be69d7d77c33e96080635a8af Mon Sep 17 00:00:00 2001 From: Andrew Gallant Date: Thu, 20 Nov 2025 14:14:53 -0500 Subject: [PATCH] [ty] Add support for relative import completions We already supported `from .. import `, but we didn't support `from ..`. This adds support for that. --- crates/ty_ide/src/completion.rs | 219 ++++++++++++++++--- crates/ty_python_semantic/src/module_name.rs | 2 +- 2 files changed, 192 insertions(+), 29 deletions(-) diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index c3ba8cf69d..0f8c029ce1 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -646,7 +646,16 @@ struct FromImport<'a> { #[derive(Clone, Debug)] enum FromImportKind { Module, - Submodule { parent: ModuleName }, + Submodule { + parent: ModuleName, + }, + Relative { + parent: ModuleName, + /// When `true`, an `import` keyword is allowed next. + /// For example, `from ...` should offer `import` + /// but also submodule completions. + import_keyword_allowed: bool, + }, Attribute, } @@ -1048,20 +1057,51 @@ impl<'a> ImportStatement<'a> { } (Some(from), import) => { let ast = find_ast_for_from_import(parsed, from)?; + // If we saw an `import` keyword, then that means the + // cursor must be *after* the `import`. And thus we + // only ever need to offer completions for importable + // elements from the module being imported. let kind = if import.is_some() { FromImportKind::Attribute - } else if initial_dot { - // TODO: Handle relative imports here. - if to_complete.starts_with('.') { - return Some(ImportStatement::Incomplete(IncompleteImport::Import)); - } - let (parent, _) = to_complete.rsplit_once('.')?; - let module_name = ModuleName::new(parent)?; - FromImportKind::Submodule { - parent: module_name, - } - } else { + } else if !initial_dot { FromImportKind::Module + } else { + let to_complete_without_leading_dots = to_complete.trim_start_matches('.'); + + // When there aren't any leading dots to trim, then we + // have a regular absolute import. Otherwise, it's relative. + if to_complete == to_complete_without_leading_dots { + let (parent, _) = to_complete.rsplit_once('.')?; + let parent = ModuleName::new(parent)?; + FromImportKind::Submodule { parent } + } else { + let all_dots = to_complete.chars().all(|c| c == '.'); + // We should suggest `import` in `from ...` + // and `from ...imp`. + let import_keyword_allowed = + all_dots || !to_complete_without_leading_dots.contains('.'); + let parent = if all_dots { + ModuleName::from_import_statement(db, file, ast).ok()? + } else { + // We know `to_complete` is not all dots. + // But that it starts with a dot. + // So we must have one of `..foo`, `..foo.` + // or `..foo.bar`. We drop the leading dots, + // since those are captured by `ast.level`. + // From there, we can treat it like a normal + // module name. We want to list submodule + // completions, so we pop off the last element + // if there are any remaining dots. + let parent = to_complete_without_leading_dots + .rsplit_once('.') + .map(|(parent, _)| parent); + ModuleName::from_identifier_parts(db, file, parent, ast.level).ok()? + }; + FromImportKind::Relative { + parent, + import_keyword_allowed, + } + } }; Some(ImportStatement::FromImport(FromImport { ast, kind })) } @@ -1093,6 +1133,15 @@ impl<'a> ImportStatement<'a> { FromImportKind::Submodule { ref parent } => { completions.extend(model.import_submodule_completions_for_name(parent)); } + FromImportKind::Relative { + ref parent, + import_keyword_allowed, + } => { + completions.extend(model.import_submodule_completions_for_name(parent)); + if import_keyword_allowed { + completions.try_add(Completion::keyword("import")); + } + } FromImportKind::Attribute => { completions.extend(model.from_import_completions(ast)); } @@ -5178,9 +5227,9 @@ match status: } #[test] - fn from_import_importt_suggests_import() { + fn from_import_importt_suggests_nothing() { let builder = completion_test_builder("from typing importt"); - assert_snapshot!(builder.build().snapshot(), @"import"); + assert_snapshot!(builder.build().snapshot(), @""); } #[test] @@ -5207,12 +5256,10 @@ match status: assert_snapshot!(builder.build().snapshot(), @"import"); } - /// The following behaviour may not be reflected in editors, since LSP - /// clients may do their own filtering of completion suggestions. #[test] - fn from_import_random_name_suggests_import() { + fn from_import_random_name_suggests_nothing() { let builder = completion_test_builder("from typing aa"); - assert_snapshot!(builder.build().snapshot(), @"import"); + assert_snapshot!(builder.build().snapshot(), @""); } #[test] @@ -5261,29 +5308,37 @@ match status: #[test] fn from_import_only_dot() { let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo.py", "") .source( - "main.py", + "package/sub1/sub2/bar.py", " - import_zqzqzq = 1 - from . - ", +import_zqzqzq = 1 +from . +", ) .completion_test_builder(); - assert_snapshot!(builder.build().snapshot(), @"import"); + assert_snapshot!(builder.build().snapshot(), @r" + import + "); } #[test] fn from_import_only_dot_incomplete() { let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo.py", "") .source( - "main.py", + "package/sub1/sub2/bar.py", " - import_zqzqzq = 1 - from .imp - ", +import_zqzqzq = 1 +from .imp +", ) .completion_test_builder(); - assert_snapshot!(builder.build().snapshot(), @"import"); + assert_snapshot!(builder.build().snapshot(), @r" + import + "); } #[test] @@ -5297,6 +5352,114 @@ match status: assert_snapshot!(builder.build().snapshot(), @"ZQZQZQ"); } + #[test] + fn relative_import_module_after_dots1() { + let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo.py", "") + .source("package/sub1/sub2/bar.py", "from ...") + .completion_test_builder(); + assert_snapshot!(builder.build().snapshot(), @r" + import + foo + "); + } + + #[test] + fn relative_import_module_after_dots2() { + let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo/__init__.py", "") + .source("package/foo/bar.py", "") + .source("package/foo/baz.py", "") + .source("package/sub1/sub2/bar.py", "from ...foo.") + .completion_test_builder(); + assert_snapshot!(builder.build().snapshot(), @r" + bar + baz + "); + } + + #[test] + fn relative_import_module_after_dots3() { + let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo.py", "") + .source("package/sub1/sub2/bar.py", "from.") + .completion_test_builder(); + assert_snapshot!(builder.build().snapshot(), @r" + import + "); + } + + #[test] + fn relative_import_module_after_dots4() { + let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo.py", "") + .source("package/sub1/bar.py", "from ..") + .completion_test_builder(); + assert_snapshot!(builder.build().snapshot(), @r" + import + foo + "); + } + + #[test] + fn relative_import_module_after_typing1() { + let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo.py", "") + .source("package/sub1/sub2/bar.py", "from ...fo") + .completion_test_builder(); + assert_snapshot!(builder.build().snapshot(), @"foo"); + } + + #[test] + fn relative_import_module_after_typing2() { + let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo/__init__.py", "") + .source("package/foo/bar.py", "") + .source("package/foo/baz.py", "") + .source("package/sub1/sub2/bar.py", "from ...foo.ba") + .completion_test_builder(); + assert_snapshot!(builder.build().snapshot(), @r" + bar + baz + "); + } + + #[test] + fn relative_import_module_after_typing3() { + let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/foo.py", "") + .source("package/imposition.py", "") + .source("package/sub1/sub2/bar.py", "from ...im") + .completion_test_builder(); + assert_snapshot!(builder.build().snapshot(), @r" + import + imposition + "); + } + + #[test] + fn relative_import_module_after_typing4() { + let builder = CursorTest::builder() + .source("package/__init__.py", "") + .source("package/sub1/__init__.py", "") + .source("package/sub1/foo.py", "") + .source("package/sub1/imposition.py", "") + .source("package/sub1/bar.py", "from ..sub1.") + .completion_test_builder(); + assert_snapshot!(builder.build().snapshot(), @r" + bar + foo + imposition + "); + } + /// A way to create a simple single-file (named `main.py`) completion test /// builder. /// diff --git a/crates/ty_python_semantic/src/module_name.rs b/crates/ty_python_semantic/src/module_name.rs index b257d0a6df..ff25d1fdbd 100644 --- a/crates/ty_python_semantic/src/module_name.rs +++ b/crates/ty_python_semantic/src/module_name.rs @@ -296,7 +296,7 @@ impl ModuleName { } /// Computes the absolute module name from the LHS components of `from LHS import RHS` - pub(crate) fn from_identifier_parts( + pub fn from_identifier_parts( db: &dyn Db, importing_file: File, module: Option<&str>,