diff --git a/pylsp/plugins/rope_autoimport.py b/pylsp/plugins/rope_autoimport.py index aaaca6b4..70a4b618 100644 --- a/pylsp/plugins/rope_autoimport.py +++ b/pylsp/plugins/rope_autoimport.py @@ -234,6 +234,14 @@ def _sort_import(score: int) -> str: return "[z" + str(score).rjust(_score_pow, "0") +def get_name_or_module(document, diagnostic) -> str: + return ( + parso.parse(document.lines[diagnostic["range"]["start"]["line"]]) + .get_leaf_for_position((1, diagnostic["range"]["start"]["character"] + 1)) + .value + ) + + @hookimpl def pylsp_code_actions( config: Config, @@ -267,10 +275,8 @@ def pylsp_code_actions( for diagnostic in context.get("diagnostics", []): if "undefined name" not in diagnostic.get("message", "").lower(): continue - expr = parso.parse(document.lines[diagnostic["range"]["start"]["line"]]) - word = expr.get_leaf_for_position( - (1, diagnostic["range"]["start"]["character"] + 1) - ).value + + word = get_name_or_module(document, diagnostic) log.debug(f"autoimport: searching for word: {word}") rope_config = config.settings(document_path=document.path).get("rope", {}) autoimport = workspace._rope_autoimport(rope_config, feature="code_actions") diff --git a/test/plugins/test_autoimport.py b/test/plugins/test_autoimport.py index 86c0d825..9aafca48 100644 --- a/test/plugins/test_autoimport.py +++ b/test/plugins/test_autoimport.py @@ -9,10 +9,14 @@ from pylsp import lsp, uris from pylsp.config.config import Config -from pylsp.plugins.rope_autoimport import _get_score, _should_insert, get_names +from pylsp.plugins.rope_autoimport import ( + _get_score, + _should_insert, + get_name_or_module, + get_names, +) from pylsp.plugins.rope_autoimport import ( pylsp_completions as pylsp_autoimport_completions, - pylsp_code_actions as pylsp_autoimport_code_actions, ) from pylsp.plugins.rope_autoimport import pylsp_initialize from pylsp.workspace import Workspace @@ -227,26 +231,20 @@ class sfa: "message", ["Undefined name `os`", "F821 undefined name 'numpy'", "undefined name 'numpy'"], ) -def test_autoimport_code_actions(config, autoimport_workspace, message): - source = "os" +def test_autoimport_code_actions_get_correct_module_name(autoimport_workspace, message): + source = "os.path.join('a', 'b')" autoimport_workspace.put_document(DOC_URI, source=source) doc = autoimport_workspace.get_document(DOC_URI) - context = { - "diagnostics": [ - { - "range": { - "start": {"line": 0, "character": 0}, - "end": {"line": 0, "character": 2}, - }, - "message": message, - } - ] + diagnostic = { + "range": { + "start": {"line": 0, "character": 0}, + "end": {"line": 0, "character": 2}, + }, + "message": message, } - actions = pylsp_autoimport_code_actions( - config, autoimport_workspace, doc, None, context - ) + module_name = get_name_or_module(doc, diagnostic) autoimport_workspace.rm_document(DOC_URI) - assert any(action.get("title") == "import os" for action in actions) + assert module_name == "os" # rope autoimport launches a sqlite database which checks from which thread it is called.