From 52f59c5c39057caa13be5a65cb51093528031bf0 Mon Sep 17 00:00:00 2001 From: Andrew Gallant Date: Mon, 1 Dec 2025 11:04:20 -0500 Subject: [PATCH] [ty] Fix auto-import code action to handle pre-existing import Previously, the code action to do auto-import on a pre-existing symbol assumed that the auto-importer would always generate an import statement. But sometimes an import statement already exists. A good example of this is the following snippet: ``` import warnings @deprecated def myfunc(): pass ``` Specifically, `deprecated` exists in `warnings` but isn't currently imported. A code action to fix this could feasibly do two transformations here. One is: ``` import warnings @warnings.deprecated def myfunc(): pass ``` Another is: ``` from warnings import deprecated import warnings @deprecated def myfunc(): pass ``` The existing auto-import infrastructure chooses the former, since it reuses a pre-existing import statement. But this PR chooses the latter for the case of a code action. I'm not 100% sure this is the correct choice, but it seems to defer more strongly to what the user has typed. That is, that they want to use it unqualified because it's what has been typed. So we should add the necessary import statement to make that work. Fixes astral-sh/ty#1668 --- crates/ty_ide/src/completion.rs | 26 ++++- crates/ty_ide/src/importer.rs | 10 ++ crates/ty_server/tests/e2e/code_actions.rs | 39 ++++++++ ...n_existing_import_undefined_decorator.snap | 98 +++++++++++++++++++ 4 files changed, 169 insertions(+), 4 deletions(-) create mode 100644 crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_existing_import_undefined_decorator.snap diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index 45576ae531..811c2f6aef 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -417,7 +417,16 @@ pub fn completion<'db>( } if settings.auto_import { if let Some(scoped) = scoped { - add_unimported_completions(db, file, &parsed, scoped, &mut completions); + add_unimported_completions( + db, + file, + &parsed, + scoped, + |module_name: &ModuleName, symbol: &str| { + ImportRequest::import_from(module_name.as_str(), symbol) + }, + &mut completions, + ); } } } @@ -453,7 +462,16 @@ pub(crate) fn missing_imports( ) -> Vec { let mut completions = Completions::exactly(db, symbol); let scoped = ScopedTarget { node }; - add_unimported_completions(db, file, parsed, scoped, &mut completions); + add_unimported_completions( + db, + file, + parsed, + scoped, + |module_name: &ModuleName, symbol: &str| { + ImportRequest::import_from(module_name.as_str(), symbol).force() + }, + &mut completions, + ); completions.into_imports() } @@ -502,6 +520,7 @@ fn add_unimported_completions<'db>( file: File, parsed: &ParsedModuleRef, scoped: ScopedTarget<'_>, + create_import_request: impl for<'a> Fn(&'a ModuleName, &'a str) -> ImportRequest<'a>, completions: &mut Completions<'db>, ) { // This is redundant since `all_symbols` will also bail @@ -523,8 +542,7 @@ fn add_unimported_completions<'db>( continue; } - let request = - ImportRequest::import_from(symbol.module.name(db).as_str(), &symbol.symbol.name); + let request = create_import_request(symbol.module.name(db), &symbol.symbol.name); // FIXME: `all_symbols` doesn't account for wildcard imports. // Since we're looking at every module, this is probably // "fine," but it might mean that we import a symbol from the diff --git a/crates/ty_ide/src/importer.rs b/crates/ty_ide/src/importer.rs index 680fb1d8cd..94b2457e74 100644 --- a/crates/ty_ide/src/importer.rs +++ b/crates/ty_ide/src/importer.rs @@ -553,6 +553,16 @@ impl<'a> ImportRequest<'a> { } } + /// Causes this request to become a command. This will force the + /// requested import style, even if another style would be more + /// appropriate generally. + pub(crate) fn force(mut self) -> Self { + Self { + force_style: true, + ..self + } + } + /// Attempts to change the import request style so that the chances /// of an import conflict are minimized (although not always reduced /// to zero). diff --git a/crates/ty_server/tests/e2e/code_actions.rs b/crates/ty_server/tests/e2e/code_actions.rs index baee6c0b42..77f4d42fcb 100644 --- a/crates/ty_server/tests/e2e/code_actions.rs +++ b/crates/ty_server/tests/e2e/code_actions.rs @@ -198,6 +198,45 @@ def my_func(): ... Ok(()) } +// Using an unimported decorator `@deprecated` +#[test] +fn code_action_existing_import_undefined_decorator() -> Result<()> { + let workspace_root = SystemPath::new("src"); + let foo = SystemPath::new("src/foo.py"); + let foo_content = r#"\ +import warnings + +@deprecated("do not use!!!") +def my_func(): ... +"#; + + let ty_toml = SystemPath::new("ty.toml"); + let ty_toml_content = ""; + + let mut server = TestServerBuilder::new()? + .with_workspace(workspace_root, None)? + .with_file(ty_toml, ty_toml_content)? + .with_file(foo, foo_content)? + .enable_pull_diagnostics(true) + .build() + .wait_until_workspaces_are_initialized(); + + server.open_text_document(foo, &foo_content, 1); + + // Wait for diagnostics to be computed. + let diagnostics = server.document_diagnostic_request(foo, None); + let range = full_range(foo_content); + let code_action_params = code_actions_at(&server, diagnostics, foo, range); + + // Get code actions + let code_action_id = server.send_request::(code_action_params); + let code_actions = server.await_response::(&code_action_id); + + insta::assert_json_snapshot!(code_actions); + + Ok(()) +} + // Accessing `typing.Literal` without `typing` imported (ideally we suggest importing `typing`) #[test] fn code_action_attribute_access_on_unimported() -> Result<()> { diff --git a/crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_existing_import_undefined_decorator.snap b/crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_existing_import_undefined_decorator.snap new file mode 100644 index 0000000000..fd022ed8b2 --- /dev/null +++ b/crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_existing_import_undefined_decorator.snap @@ -0,0 +1,98 @@ +--- +source: crates/ty_server/tests/e2e/code_actions.rs +expression: code_actions +--- +[ + { + "title": "import warnings.deprecated", + "kind": "quickfix", + "diagnostics": [ + { + "range": { + "start": { + "line": 3, + "character": 1 + }, + "end": { + "line": 3, + "character": 11 + } + }, + "severity": 1, + "code": "unresolved-reference", + "codeDescription": { + "href": "https://ty.dev/rules#unresolved-reference" + }, + "source": "ty", + "message": "Name `deprecated` used when not defined", + "relatedInformation": [] + } + ], + "edit": { + "changes": { + "file:///src/foo.py": [ + { + "range": { + "start": { + "line": 0, + "character": 0 + }, + "end": { + "line": 0, + "character": 0 + } + }, + "newText": "from warnings import deprecated\n" + } + ] + } + }, + "isPreferred": true + }, + { + "title": "Ignore 'unresolved-reference' for this line", + "kind": "quickfix", + "diagnostics": [ + { + "range": { + "start": { + "line": 3, + "character": 1 + }, + "end": { + "line": 3, + "character": 11 + } + }, + "severity": 1, + "code": "unresolved-reference", + "codeDescription": { + "href": "https://ty.dev/rules#unresolved-reference" + }, + "source": "ty", + "message": "Name `deprecated` used when not defined", + "relatedInformation": [] + } + ], + "edit": { + "changes": { + "file:///src/foo.py": [ + { + "range": { + "start": { + "line": 3, + "character": 28 + }, + "end": { + "line": 3, + "character": 28 + } + }, + "newText": " # ty:ignore[unresolved-reference]" + } + ] + } + }, + "isPreferred": false + } +]