Repository: python-lsp/python-lsp-server (python)

Issue:
Add code action for implementing auto-import
Hi,
while the auto-import feature is great, it is not enough. And a code action to do import of a certain class from module could be a very good improvement.  

Pull request:
Add code completions to `rope_autoimport` plugin
Fixes #403 

**Testing**

- [x] Unit test
- [x] Manual test
![image](https://github.com/python-lsp/python-lsp-server/assets/91616041/5141988d-476e-4220-82a6-f07ed4768710)


Patch:
diff --git a/CONFIGURATION.md b/CONFIGURATION.md
index f88e425..acf8a85 100644
--- a/CONFIGURATION.md
+++ b/CONFIGURATION.md
@@ -64,7 +64,9 @@ This server can be configured using the `workspace/didChangeConfiguration` metho
 | `pylsp.plugins.pylint.enabled` | `boolean` | Enable or disable the plugin. | `false` |
 | `pylsp.plugins.pylint.args` | `array` of non-unique `string` items | Arguments to pass to pylint. | `[]` |
 | `pylsp.plugins.pylint.executable` | `string` | Executable to run pylint with. Enabling this will run pylint on unsaved files via stdin. Can slow down workflow. Only works with python3. | `null` |
-| `pylsp.plugins.rope_autoimport.enabled` | `boolean` | Enable or disable autoimport. | `false` |
+| `pylsp.plugins.rope_autoimport.enabled` | `boolean` | Enable or disable autoimport. If false, neither completions nor code actions are enabled. If true, the respective features can be enabled or disabled individually. | `false` |
+| `pylsp.plugins.rope_autoimport.completions.enabled` | `boolean` | Enable or disable autoimport completions. | `true` |
+| `pylsp.plugins.rope_autoimport.code_actions.enabled` | `boolean` | Enable or disable autoimport code actions (e.g. for quick fixes). | `true` |
 | `pylsp.plugins.rope_autoimport.memory` | `boolean` | Make the autoimport database memory only. Drastically increases startup time. | `false` |
 | `pylsp.plugins.rope_completion.enabled` | `boolean` | Enable or disable the plugin. | `false` |
 | `pylsp.plugins.rope_completion.eager` | `boolean` | Resolve documentation and detail eagerly. | `false` |
diff --git a/docs/autoimport.md b/docs/autoimport.md
index 5bf573e..893a5e9 100644
--- a/docs/autoimport.md
+++ b/docs/autoimport.md
@@ -4,6 +4,7 @@ Requirements:
 
 1. install `python-lsp-server[rope]`
 2. set `pylsp.plugins.rope_autoimport.enabled` to `true`
+3. This enables both completions and code actions. You can switch them off by setting `pylsp.plugins.rope_autoimport.completions.enabled` and/or `pylsp.plugins.rope_autoimport.code_actions.enabled` to `false`
 
 ## Startup
 
diff --git a/pylsp/config/schema.json b/pylsp/config/schema.json
index fbf7f01..ba1d36f 100644
--- a/pylsp/config/schema.json
+++ b/pylsp/config/schema.json
@@ -6,11 +6,16 @@
   "properties": {
     "pylsp.configurationSources": {
       "type": "array",
-      "default": ["pycodestyle"],
+      "default": [
+        "pycodestyle"
+      ],
       "description": "List of configuration sources to use.",
       "items": {
         "type": "string",
-        "enum": ["pycodestyle", "flake8"]
+        "enum": [
+          "pycodestyle",
+          "flake8"
+        ]
       },
       "uniqueItems": true
     },
@@ -20,7 +25,10 @@
       "description": "Enable or disable the plugin (disabling required to use `yapf`)."
     },
     "pylsp.plugins.flake8.config": {
-      "type": ["string", "null"],
+      "type": [
+        "string",
+        "null"
+      ],
       "default": null,
       "description": "Path to the config file that will be the authoritative config source."
     },
@@ -51,12 +59,18 @@
       "description": "Path to the flake8 executable."
     },
     "pylsp.plugins.flake8.filename": {
-      "type": ["string", "null"],
+      "type": [
+        "string",
+        "null"
+      ],
       "default": null,
       "description": "Only check for filenames matching the patterns in this list."
     },
     "pylsp.plugins.flake8.hangClosing": {
-      "type": ["boolean", "null"],
+      "type": [
+        "boolean",
+        "null"
+      ],
       "default": null,
       "description": "Hang closing bracket instead of matching indentation of opening bracket's line."
     },
@@ -74,17 +88,25 @@
       "description": "Maximum allowed complexity threshold."
     },
     "pylsp.plugins.flake8.maxLineLength": {
-      "type": ["integer", "null"],
+      "type": [
+        "integer",
+        "null"
+      ],
       "default": null,
       "description": "Maximum allowed line length for the entirety of this run."
     },
     "pylsp.plugins.flake8.indentSize": {
-      "type": ["integer", "null"],
+      "type": [
+        "integer",
+        "null"
+      ],
       "default": null,
       "description": "Set indentation spaces."
     },
     "pylsp.plugins.flake8.perFileIgnores": {
-      "type": ["array"],
+      "type": [
+        "array"
+      ],
       "default": [],
       "items": {
         "type": "string"
@@ -92,7 +114,10 @@
       "description": "A pairing of filenames and violation codes that defines which violations to ignore in a particular file, for example: `[\"file_path.py:W305,W304\"]`)."
     },
     "pylsp.plugins.flake8.select": {
-      "type": ["array", "null"],
+      "type": [
+        "array",
+        "null"
+      ],
       "default": null,
       "items": {
         "type": "string"
@@ -102,7 +127,9 @@
     },
     "pylsp.plugins.jedi.auto_import_modules": {
       "type": "array",
-      "default": ["numpy"],
+      "default": [
+        "numpy"
+      ],
       "items": {
         "type": "string"
       },
@@ -117,12 +144,18 @@
       "description": "Define extra paths for jedi.Script."
     },
     "pylsp.plugins.jedi.env_vars": {
-      "type": ["object", "null"],
+      "type": [
+        "object",
+        "null"
+      ],
       "default": null,
       "description": "Define environment variables for jedi.Script and Jedi.names."
     },
     "pylsp.plugins.jedi.environment": {
-      "type": ["string", "null"],
+      "type": [
+        "string",
+        "null"
+      ],
       "default": null,
       "description": "Define environment for jedi.Script and Jedi.names."
     },
@@ -166,7 +199,12 @@
       "items": {
         "type": "string"
       },
-      "default": ["pandas", "numpy", "tensorflow", "matplotlib"],
+      "default": [
+        "pandas",
+        "numpy",
+        "tensorflow",
+        "matplotlib"
+      ],
       "description": "Modules for which labels and snippets should be cached."
     },
     "pylsp.plugins.jedi_definition.enabled": {
@@ -267,7 +305,10 @@
       "description": "When parsing directories, only check filenames matching these patterns."
     },
     "pylsp.plugins.pycodestyle.select": {
-      "type": ["array", "null"],
+      "type": [
+        "array",
+        "null"
+      ],
       "default": null,
       "items": {
         "type": "string"
@@ -285,17 +326,26 @@
       "description": "Ignore errors and warnings"
     },
     "pylsp.plugins.pycodestyle.hangClosing": {
-      "type": ["boolean", "null"],
+      "type": [
+        "boolean",
+        "null"
+      ],
       "default": null,
       "description": "Hang closing bracket instead of matching indentation of opening bracket's line."
     },
     "pylsp.plugins.pycodestyle.maxLineLength": {
-      "type": ["integer", "null"],
+      "type": [
+        "integer",
+        "null"
+      ],
       "default": null,
       "description": "Set maximum allowed line length."
     },
     "pylsp.plugins.pycodestyle.indentSize": {
-      "type": ["integer", "null"],
+      "type": [
+        "integer",
+        "null"
+      ],
       "default": null,
       "description": "Set indentation spaces."
     },
@@ -305,9 +355,17 @@
       "description": "Enable or disable the plugin."
     },
     "pylsp.plugins.pydocstyle.convention": {
-      "type": ["string", "null"],
+      "type": [
+        "string",
+        "null"
+      ],
       "default": null,
-      "enum": ["pep257", "numpy", "google", null],
+      "enum": [
+        "pep257",
+        "numpy",
+        "google",
+        null
+      ],
       "description": "Choose the basic list of checked errors by specifying an existing convention."
     },
     "pylsp.plugins.pydocstyle.addIgnore": {
@@ -338,7 +396,10 @@
       "description": "Ignore errors and warnings"
     },
     "pylsp.plugins.pydocstyle.select": {
-      "type": ["array", "null"],
+      "type": [
+        "array",
+        "null"
+      ],
       "default": null,
       "items": {
         "type": "string"
@@ -376,14 +437,27 @@
       "description": "Arguments to pass to pylint."
     },
     "pylsp.plugins.pylint.executable": {
-      "type": ["string", "null"],
+      "type": [
+        "string",
+        "null"
+      ],
       "default": null,
       "description": "Executable to run pylint with. Enabling this will run pylint on unsaved files via stdin. Can slow down workflow. Only works with python3."
     },
     "pylsp.plugins.rope_autoimport.enabled": {
       "type": "boolean",
       "default": false,
-      "description": "Enable or disable autoimport."
+      "description": "Enable or disable autoimport. If false, neither completions nor code actions are enabled. If true, the respective features can be enabled or disabled individually."
+    },
+    "pylsp.plugins.rope_autoimport.completions.enabled": {
+      "type": "boolean",
+      "default": true,
+      "description": "Enable or disable autoimport completions."
+    },
+    "pylsp.plugins.rope_autoimport.code_actions.enabled": {
+      "type": "boolean",
+      "default": true,
+      "description": "Enable or disable autoimport code actions (e.g. for quick fixes)."
     },
     "pylsp.plugins.rope_autoimport.memory": {
       "type": "boolean",
@@ -406,12 +480,18 @@
       "description": "Enable or disable the plugin."
     },
     "pylsp.rope.extensionModules": {
-      "type": ["string", "null"],
+      "type": [
+        "string",
+        "null"
+      ],
       "default": null,
       "description": "Builtin and c-extension modules that are allowed to be imported and inspected by rope."
     },
     "pylsp.rope.ropeFolder": {
-      "type": ["array", "null"],
+      "type": [
+        "array",
+        "null"
+      ],
       "default": null,
       "items": {
         "type": "string"
@@ -420,4 +500,4 @@
       "description": "The name of the folder in which rope stores project configurations and data.  Pass `null` for not using such a folder at all."
     }
   }
-}
+}
\ No newline at end of file
diff --git a/pylsp/plugins/rope_autoimport.py b/pylsp/plugins/rope_autoimport.py
index 1caab35..ca3db1c 100644
--- a/pylsp/plugins/rope_autoimport.py
+++ b/pylsp/plugins/rope_autoimport.py
@@ -21,13 +21,27 @@ log = logging.getLogger(__name__)
 
 _score_pow = 5
 _score_max = 10**_score_pow
-MAX_RESULTS = 1000
+MAX_RESULTS_COMPLETIONS = 1000
+MAX_RESULTS_CODE_ACTIONS = 5
 
 
 @hookimpl
 def pylsp_settings() -> Dict[str, Dict[str, Dict[str, Any]]]:
     # Default rope_completion to disabled
-    return {"plugins": {"rope_autoimport": {"enabled": False, "memory": False}}}
+    return {
+        "plugins": {
+            "rope_autoimport": {
+                "enabled": False,
+                "memory": False,
+                "completions": {
+                    "enabled": True,
+                },
+                "code_actions": {
+                    "enabled": True,
+                },
+            }
+        }
+    }
 
 
 # pylint: disable=too-many-return-statements
@@ -122,6 +136,7 @@ def _process_statements(
     word: str,
     autoimport: AutoImport,
     document: Document,
+    feature: str = "completions",
 ) -> Generator[Dict[str, Any], None, None]:
     for suggestion in suggestions:
         insert_line = autoimport.find_insertion_line(document.source) - 1
@@ -134,14 +149,26 @@ def _process_statements(
         if score > _score_max:
             continue
         # TODO make this markdown
-        yield {
-            "label": suggestion.name,
-            "kind": suggestion.itemkind,
-            "sortText": _sort_import(score),
-            "data": {"doc_uri": doc_uri},
-            "detail": _document(suggestion.import_statement),
-            "additionalTextEdits": [edit],
-        }
+        if feature == "completions":
+            yield {
+                "label": suggestion.name,
+                "kind": suggestion.itemkind,
+                "sortText": _sort_import(score),
+                "data": {"doc_uri": doc_uri},
+                "detail": _document(suggestion.import_statement),
+                "additionalTextEdits": [edit],
+            }
+        elif feature == "code_actions":
+            yield {
+                "title": suggestion.import_statement,
+                "kind": "quickfix",
+                "edit": {"changes": {doc_uri: [edit]}},
+                # data is a supported field for codeAction responses
+                # See https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_codeAction  # pylint: disable=line-too-long
+                "data": {"sortText": _sort_import(score)},
+            }
+        else:
+            raise ValueError(f"Unknown feature: {feature}")
 
 
 def get_names(script: Script) -> Set[str]:
@@ -160,6 +187,13 @@ def pylsp_completions(
     ignored_names: Union[Set[str], None],
 ):
     """Get autoimport suggestions."""
+    if (
+        not config.plugin_settings("rope_autoimport")
+        .get("completions", {})
+        .get("enabled", True)
+    ):
+        return []
+
     line = document.lines[position["line"]]
     expr = parso.parse(line)
     word_node = expr.get_leaf_for_position((1, position["character"]))
@@ -175,12 +209,14 @@ def pylsp_completions(
     suggestions = list(autoimport.search_full(word, ignored_names=ignored_names))
     results = list(
         sorted(
-            _process_statements(suggestions, document.uri, word, autoimport, document),
+            _process_statements(
+                suggestions, document.uri, word, autoimport, document, "completions"
+            ),
             key=lambda statement: statement["sortText"],
         )
     )
-    if len(results) > MAX_RESULTS:
-        results = results[:MAX_RESULTS]
+    if len(results) > MAX_RESULTS_COMPLETIONS:
+        results = results[:MAX_RESULTS_COMPLETIONS]
     return results
 
 
@@ -206,6 +242,83 @@ def _sort_import(score: int) -> str:
     return "[z" + str(score).rjust(_score_pow, "0")
 
 
+def get_name_or_module(document, diagnostic) -> str:
+    start = diagnostic["range"]["start"]
+    return (
+        parso.parse(document.lines[start["line"]])
+        .get_leaf_for_position((1, start["character"] + 1))
+        .value
+    )
+
+
+@hookimpl
+def pylsp_code_actions(
+    config: Config,
+    workspace: Workspace,
+    document: Document,
+    range: Dict,  # pylint: disable=redefined-builtin
+    context: Dict,
+) -> List[Dict]:
+    """
+    Provide code actions through rope.
+
+    Parameters
+    ----------
+    config : pylsp.config.config.Config
+        Current config.
+    workspace : pylsp.workspace.Workspace
+        Current workspace.
+    document : pylsp.workspace.Document
+        Document to apply code actions on.
+    range : Dict
+        Range argument given by pylsp. Not used here.
+    context : Dict
+        CodeActionContext given as dict.
+
+    Returns
+    -------
+      List of dicts containing the code actions.
+    """
+    if (
+        not config.plugin_settings("rope_autoimport")
+        .get("code_actions", {})
+        .get("enabled", True)
+    ):
+        return []
+
+    log.debug(f"textDocument/codeAction: {document} {range} {context}")
+    code_actions = []
+    for diagnostic in context.get("diagnostics", []):
+        if "undefined name" not in diagnostic.get("message", "").lower():
+            continue
+
+        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")
+        suggestions = list(autoimport.search_full(word))
+        log.debug("autoimport: suggestions: %s", suggestions)
+        results = list(
+            sorted(
+                _process_statements(
+                    suggestions,
+                    document.uri,
+                    word,
+                    autoimport,
+                    document,
+                    "code_actions",
+                ),
+                key=lambda statement: statement["data"]["sortText"],
+            )
+        )
+
+        if len(results) > MAX_RESULTS_CODE_ACTIONS:
+            results = results[:MAX_RESULTS_CODE_ACTIONS]
+        code_actions.extend(results)
+
+    return code_actions
+
+
 def _reload_cache(
     config: Config, workspace: Workspace, files: Optional[List[Document]] = None
 ):
diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
index 760ad97..52a22a3 100644
--- a/pylsp/python_lsp.py
+++ b/pylsp/python_lsp.py
@@ -385,7 +385,7 @@ class PythonLSPServer(MethodDispatcher):
     def m_initialized(self, **_kwargs):
         self._hook("pylsp_initialized")
 
-    def code_actions(self, doc_uri, range, context):
+    def code_actions(self, doc_uri: str, range: Dict, context: Dict):
         return flatten(
             self._hook("pylsp_code_actions", doc_uri, range=range, context=context)
         )
diff --git a/pylsp/workspace.py b/pylsp/workspace.py
index fb524c7..5c6880c 100644
--- a/pylsp/workspace.py
+++ b/pylsp/workspace.py
@@ -8,7 +8,7 @@ import os
 import re
 import uuid
 import functools
-from typing import Optional, Generator, Callable, List
+from typing import Literal, Optional, Generator, Callable, List
 from threading import RLock
 
 import jedi
@@ -58,16 +58,30 @@ class Workspace:
         # Whilst incubating, keep rope private
         self.__rope = None
         self.__rope_config = None
-        self.__rope_autoimport = None
 
-    def _rope_autoimport(self, rope_config: Optional, memory: bool = False):
+        # We have a sperate AutoImport object for each feature to avoid sqlite errors
+        # from accessing the same database from multiple threads.
+        # An upstream fix discussion is here: https://github.com/python-rope/rope/issues/713
+        self.__rope_autoimport = (
+            {}
+        )  # Type: Dict[Literal["completions", "code_actions"], rope.contrib.autoimport.sqlite.AutoImport]
+
+    def _rope_autoimport(
+        self,
+        rope_config: Optional,
+        memory: bool = False,
+        feature: Literal["completions", "code_actions"] = "completions",
+    ):
         # pylint: disable=import-outside-toplevel
         from rope.contrib.autoimport.sqlite import AutoImport
 
-        if self.__rope_autoimport is None:
+        if feature not in ["completions", "code_actions"]:
+            raise ValueError(f"Unknown feature {feature}")
+
+        if self.__rope_autoimport.get(feature, None) is None:
             project = self._rope_project_builder(rope_config)
-            self.__rope_autoimport = AutoImport(project, memory=memory)
-        return self.__rope_autoimport
+            self.__rope_autoimport[feature] = AutoImport(project, memory=memory)
+        return self.__rope_autoimport[feature]
 
     def _rope_project_builder(self, rope_config):
         # pylint: disable=import-outside-toplevel
@@ -374,8 +388,8 @@ class Workspace:
         )
 
     def close(self):
-        if self.__rope_autoimport is not None:
-            self.__rope_autoimport.close()
+        for _, autoimport in self.__rope_autoimport.items():
+            autoimport.close()
 
 
 class Document:


Test patch:
diff --git a/test/fixtures.py b/test/fixtures.py
index ed6206a..11c302b 100644
--- a/test/fixtures.py
+++ b/test/fixtures.py
@@ -8,7 +8,6 @@ from unittest.mock import MagicMock
 from test.test_utils import ClientServerPair, CALL_TIMEOUT_IN_SECONDS
 
 import pytest
-import pylsp_jsonrpc
 
 from pylsp_jsonrpc.dispatchers import MethodDispatcher
 from pylsp_jsonrpc.endpoint import Endpoint
@@ -176,13 +175,8 @@ def client_server_pair():
 
     yield (client_server_pair_obj.client, client_server_pair_obj.server)
 
-    try:
-        shutdown_response = client_server_pair_obj.client._endpoint.request(
-            "shutdown"
-        ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
-        assert shutdown_response is None
-        client_server_pair_obj.client._endpoint.notify("exit")
-    except pylsp_jsonrpc.exceptions.JsonRpcInvalidParams:
-        # SQLite objects created in a thread can only be used in that same thread.
-        # This exeception is raised when testing rope autoimport.
-        client_server_pair_obj.client._endpoint.notify("exit")
+    shutdown_response = client_server_pair_obj.client._endpoint.request(
+        "shutdown"
+    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    assert shutdown_response is None
+    client_server_pair_obj.client._endpoint.notify("exit")
diff --git a/test/plugins/test_autoimport.py b/test/plugins/test_autoimport.py
index b1c4677..ec5c0a3 100644
--- a/test/plugins/test_autoimport.py
+++ b/test/plugins/test_autoimport.py
@@ -1,18 +1,20 @@
 # Copyright 2022- Python Language Server Contributors.
 
 from typing import Any, Dict, List
-from unittest.mock import Mock, patch
-
-from test.test_notebook_document import wait_for_condition
-from test.test_utils import send_initialize_request, send_notebook_did_open
+from unittest.mock import Mock
 
 import jedi
 import parso
 import pytest
 
-from pylsp import IS_WIN, lsp, uris
+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,
 )
@@ -37,7 +39,16 @@ def autoimport_workspace(tmp_path_factory) -> Workspace:
         uris.from_fs_path(str(tmp_path_factory.mktemp("pylsp"))), Mock()
     )
     workspace._config = Config(workspace.root_uri, {}, 0, {})
-    workspace._config.update({"rope_autoimport": {"memory": True, "enabled": True}})
+    workspace._config.update(
+        {
+            "rope_autoimport": {
+                "memory": True,
+                "enabled": True,
+                "completions": {"enabled": True},
+                "code_actions": {"enabled": True},
+            }
+        }
+    )
     pylsp_initialize(workspace._config, workspace)
     yield workspace
     workspace.close()
@@ -216,72 +227,103 @@ def test_get_names():
     assert results == set(["blah", "bleh", "e", "hello", "someone", "sfa", "a", "b"])
 
 
-@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
-def test_autoimport_for_notebook_document(
-    client_server_pair,
-):
-    client, server = client_server_pair
-    send_initialize_request(client)
-
-    with patch.object(server._endpoint, "notify") as mock_notify:
-        # Expectations:
-        # 1. We receive an autoimport suggestion for "os" in the first cell because
-        #    os is imported after that.
-        # 2. We don't receive an autoimport suggestion for "os" in the second cell because it's
-        #    already imported in the second cell.
-        # 3. We don't receive an autoimport suggestion for "os" in the third cell because it's
-        #    already imported in the second cell.
-        # 4. We receive an autoimport suggestion for "sys" because it's not already imported
-        send_notebook_did_open(client, ["os", "import os\nos", "os", "sys"])
-        wait_for_condition(lambda: mock_notify.call_count >= 3)
-
-    server.m_workspace__did_change_configuration(
-        settings={
-            "pylsp": {"plugins": {"rope_autoimport": {"enabled": True, "memory": True}}}
-        }
-    )
-    rope_autoimport_settings = server.workspace._config.plugin_settings(
-        "rope_autoimport"
-    )
-    assert rope_autoimport_settings.get("enabled", False) is True
-    assert rope_autoimport_settings.get("memory", False) is True
-
-    # 1.
-    suggestions = server.completions("cell_1_uri", {"line": 0, "character": 2}).get(
-        "items"
-    )
-    assert any(
-        suggestion
-        for suggestion in suggestions
-        if contains_autoimport(suggestion, "os")
-    )
-
-    # 2.
-    suggestions = server.completions("cell_2_uri", {"line": 1, "character": 2}).get(
-        "items"
-    )
-    assert not any(
-        suggestion
-        for suggestion in suggestions
-        if contains_autoimport(suggestion, "os")
-    )
-
-    # 3.
-    suggestions = server.completions("cell_3_uri", {"line": 0, "character": 2}).get(
-        "items"
-    )
-    assert not any(
-        suggestion
-        for suggestion in suggestions
-        if contains_autoimport(suggestion, "os")
-    )
-
-    # 4.
-    suggestions = server.completions("cell_4_uri", {"line": 0, "character": 3}).get(
-        "items"
-    )
-    assert any(
-        suggestion
-        for suggestion in suggestions
-        if contains_autoimport(suggestion, "sys")
-    )
+# Tests ruff, flake8 and pyflakes messages
+@pytest.mark.parametrize(
+    "message",
+    ["Undefined name `os`", "F821 undefined name 'numpy'", "undefined name 'numpy'"],
+)
+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)
+    diagnostic = {
+        "range": {
+            "start": {"line": 0, "character": 0},
+            "end": {"line": 0, "character": 2},
+        },
+        "message": message,
+    }
+    module_name = get_name_or_module(doc, diagnostic)
+    autoimport_workspace.rm_document(DOC_URI)
+    assert module_name == "os"
+
+
+# rope autoimport launches a sqlite database which checks from which thread it is called.
+# This makes the test below fail because we access the db from a different thread.
+# See https://stackoverflow.com/questions/48218065/objects-created-in-a-thread-can-only-be-used-in-that-same-thread
+# @pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+# def test_autoimport_completions_for_notebook_document(
+#     client_server_pair,
+# ):
+#     client, server = client_server_pair
+#     send_initialize_request(client)
+
+#     with patch.object(server._endpoint, "notify") as mock_notify:
+#         # Expectations:
+#         # 1. We receive an autoimport suggestion for "os" in the first cell because
+#         #    os is imported after that.
+#         # 2. We don't receive an autoimport suggestion for "os" in the second cell because it's
+#         #    already imported in the second cell.
+#         # 3. We don't receive an autoimport suggestion for "os" in the third cell because it's
+#         #    already imported in the second cell.
+#         # 4. We receive an autoimport suggestion for "sys" because it's not already imported
+#         send_notebook_did_open(client, ["os", "import os\nos", "os", "sys"])
+#         wait_for_condition(lambda: mock_notify.call_count >= 3)
+
+#     server.m_workspace__did_change_configuration(
+#         settings={
+#             "pylsp": {
+#                 "plugins": {
+#                     "rope_autoimport": {
+#                         "memory": True,
+#                         "completions": {"enabled": True},
+#                     },
+#                 }
+#             }
+#         }
+#     )
+#     rope_autoimport_settings = server.workspace._config.plugin_settings(
+#         "rope_autoimport"
+#     )
+#     assert rope_autoimport_settings.get("completions", {}).get("enabled", False) is True
+#     assert rope_autoimport_settings.get("memory", False) is True
+
+#     # 1.
+#     suggestions = server.completions("cell_1_uri", {"line": 0, "character": 2}).get(
+#         "items"
+#     )
+#     assert any(
+#         suggestion
+#         for suggestion in suggestions
+#         if contains_autoimport(suggestion, "os")
+#     )
+
+#     # 2.
+#     suggestions = server.completions("cell_2_uri", {"line": 1, "character": 2}).get(
+#         "items"
+#     )
+#     assert not any(
+#         suggestion
+#         for suggestion in suggestions
+#         if contains_autoimport(suggestion, "os")
+#     )
+
+#     # 3.
+#     suggestions = server.completions("cell_3_uri", {"line": 0, "character": 2}).get(
+#         "items"
+#     )
+#     assert not any(
+#         suggestion
+#         for suggestion in suggestions
+#         if contains_autoimport(suggestion, "os")
+#     )
+
+#     # 4.
+#     suggestions = server.completions("cell_4_uri", {"line": 0, "character": 3}).get(
+#         "items"
+#     )
+#     assert any(
+#         suggestion
+#         for suggestion in suggestions
+#         if contains_autoimport(suggestion, "sys")
+#     )