[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
This commit is contained in:
Andrew Gallant 2025-12-01 11:04:20 -05:00 committed by Andrew Gallant
parent 53299cbff4
commit 52f59c5c39
4 changed files with 169 additions and 4 deletions

View File

@ -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<ImportEdit> {
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

View File

@ -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).

View File

@ -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::<CodeActionRequest>(code_action_params);
let code_actions = server.await_response::<CodeActionRequest>(&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<()> {

View File

@ -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://<temp_dir>/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://<temp_dir>/src/foo.py": [
{
"range": {
"start": {
"line": 3,
"character": 28
},
"end": {
"line": 3,
"character": 28
}
},
"newText": " # ty:ignore[unresolved-reference]"
}
]
}
},
"isPreferred": false
}
]