Ignore notebook names on cell completion for autoimport (#466)

Author: tkrabel-db

pylsp/hookspecs.py

@@ -25,7 +25,7 @@ def pylsp_commands(config, workspace):
 
 
 @hookspec
-def pylsp_completions(config, workspace, document, position):
+def pylsp_completions(config, workspace, document, position, ignored_names):
     pass
 
 

pylsp/plugins/rope_autoimport.py

@@ -1,7 +1,7 @@
 # Copyright 2022- Python Language Server Contributors.
 
 import logging
-from typing import Any, Dict, Generator, List, Optional, Set
+from typing import Any, Dict, Generator, List, Optional, Set, Union
 
 import parso
 from jedi import Script
@@ -153,7 +153,11 @@ def get_names(script: Script) -> Set[str]:
 
 @hookimpl
 def pylsp_completions(
-    config: Config, workspace: Workspace, document: Document, position
+    config: Config,
+    workspace: Workspace,
+    document: Document,
+    position,
+    ignored_names: Union[Set[str], None],
 ):
     """Get autoimport suggestions."""
     line = document.lines[position["line"]]
@@ -164,7 +168,9 @@ def pylsp_completions(
     word = word_node.value
     log.debug(f"autoimport: searching for word: {word}")
     rope_config = config.settings(document_path=document.path).get("rope", {})
-    ignored_names: Set[str] = get_names(document.jedi_script(use_document_path=True))
+    ignored_names: Set[str] = ignored_names or get_names(
+        document.jedi_script(use_document_path=True)
+    )
     autoimport = workspace._rope_autoimport(rope_config)
     suggestions = list(autoimport.search_full(word, ignored_names=ignored_names))
     results = list(

pylsp/python_lsp.py

@@ -394,7 +394,16 @@ def code_lens(self, doc_uri):
         return flatten(self._hook("pylsp_code_lens", doc_uri))
 
     def completions(self, doc_uri, position):
-        completions = self._hook("pylsp_completions", doc_uri, position=position)
+        workspace = self._match_uri_to_workspace(doc_uri)
+        document = workspace.get_document(doc_uri)
+        ignored_names = None
+        if isinstance(document, Cell):
+            # We need to get the ignored names from the whole notebook document
+            notebook_document = workspace.get_maybe_document(document.notebook_uri)
+            ignored_names = notebook_document.jedi_names(doc_uri)
+        completions = self._hook(
+            "pylsp_completions", doc_uri, position=position, ignored_names=ignored_names
+        )
         return {"isIncomplete": False, "items": flatten(completions)}
 
     def completion_item_resolve(self, completion_item):

pylsp/workspace.py

@@ -595,6 +595,7 @@ def __init__(
         self.version = version
         self.cells = cells or []
         self.metadata = metadata or {}
+        self._lock = RLock()
 
     def __str__(self):
         return "Notebook with URI '%s'" % str(self.uri)
@@ -625,6 +626,31 @@ def cell_data(self):
             offset += num_lines
         return cell_data
 
+    @lock
+    def jedi_names(
+        self,
+        up_to_cell_uri: Optional[str] = None,
+        all_scopes=False,
+        definitions=True,
+        references=False,
+    ):
+        """
+        Get the names in the notebook up to a certain cell.
+
+        Parameters
+        ----------
+        up_to_cell_uri: str, optional
+            The cell uri to stop at. If None, all cells are considered.
+        """
+        names = set()
+        for cell in self.cells:
+            cell_uri = cell["document"]
+            cell_document = self.workspace.get_cell_document(cell_uri)
+            names.update(cell_document.jedi_names(all_scopes, definitions, references))
+            if cell_uri == up_to_cell_uri:
+                break
+        return set(name.name for name in names)
+
 
 class Cell(Document):
     """

test/fixtures.py

@@ -5,9 +5,11 @@
 from io import StringIO
 from unittest.mock import MagicMock
 
-from test.test_utils import ClientServerPair
+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
 from pylsp_jsonrpc.exceptions import JsonRpcException
@@ -24,7 +26,6 @@
 def main():
     print sys.stdin.read()
 """
-CALL_TIMEOUT_IN_SECONDS = 30
 
 
 class FakeEditorMethodsMixin:
@@ -175,8 +176,13 @@ def client_server_pair():
 
     yield (client_server_pair_obj.client, client_server_pair_obj.server)
 
-    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")
+    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")

test/plugins/test_autoimport.py

@@ -1,13 +1,16 @@
 # Copyright 2022- Python Language Server Contributors.
 
-from typing import Dict, List
-from unittest.mock import Mock
+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
 
 import jedi
 import parso
 import pytest
 
-from pylsp import lsp, uris
+from pylsp import IS_WIN, 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 (
@@ -16,9 +19,17 @@
 from pylsp.plugins.rope_autoimport import pylsp_initialize
 from pylsp.workspace import Workspace
 
+
 DOC_URI = uris.from_fs_path(__file__)
 
 
+def contains_autoimport(suggestion: Dict[str, Any], module: str) -> bool:
+    """Checks if `suggestion` contains an autoimport for `module`."""
+    return suggestion.get("label", "") == module and "import" in suggestion.get(
+        "detail", ""
+    )
+
+
 @pytest.fixture(scope="session")
 def autoimport_workspace(tmp_path_factory) -> Workspace:
     "Special autoimport workspace. Persists across sessions to make in-memory sqlite3 database fast."
@@ -39,7 +50,9 @@ def completions(config: Config, autoimport_workspace: Workspace, request):
     com_position = {"line": 0, "character": position}
     autoimport_workspace.put_document(DOC_URI, source=document)
     doc = autoimport_workspace.get_document(DOC_URI)
-    yield pylsp_autoimport_completions(config, autoimport_workspace, doc, com_position)
+    yield pylsp_autoimport_completions(
+        config, autoimport_workspace, doc, com_position, None
+    )
     autoimport_workspace.rm_document(DOC_URI)
 
 
@@ -141,45 +154,13 @@ def test_autoimport_defined_name(config, workspace):
     com_position = {"line": 1, "character": 3}
     workspace.put_document(DOC_URI, source=document)
     doc = workspace.get_document(DOC_URI)
-    completions = pylsp_autoimport_completions(config, workspace, doc, com_position)
+    completions = pylsp_autoimport_completions(
+        config, workspace, doc, com_position, None
+    )
     workspace.rm_document(DOC_URI)
     assert not check_dict({"label": "List"}, completions)
 
 
-# This test has several large issues.
-# 1. autoimport relies on its sources being written to disk. This makes testing harder
-# 2. the hook doesn't handle removed files
-# 3. The testing framework cannot access the actual autoimport object so it cannot clear the cache
-# def test_autoimport_update_module(config: Config, workspace: Workspace):
-#     document2 = "SomethingYouShouldntWrite = 1"
-#     document = """SomethingYouShouldntWrit"""
-#     com_position = {
-#         "line": 0,
-#         "character": 3,
-#     }
-#     doc2_path = workspace.root_path + "/test_file_no_one_should_write_to.py"
-#     if os.path.exists(doc2_path):
-#         os.remove(doc2_path)
-#     DOC2_URI = uris.from_fs_path(doc2_path)
-#     workspace.put_document(DOC_URI, source=document)
-#     doc = workspace.get_document(DOC_URI)
-#     completions = pylsp_autoimport_completions(config, workspace, doc, com_position)
-#     assert len(completions) == 0
-#     with open(doc2_path, "w") as f:
-#         f.write(document2)
-#     workspace.put_document(DOC2_URI, source=document2)
-#     doc2 = workspace.get_document(DOC2_URI)
-#     pylsp_document_did_save(config, workspace, doc2)
-#     assert check_dict({"label": "SomethingYouShouldntWrite"}, completions)
-#     workspace.put_document(DOC2_URI, source="\n")
-#     doc2 = workspace.get_document(DOC2_URI)
-#     os.remove(doc2_path)
-#     pylsp_document_did_save(config, workspace, doc2)
-#     completions = pylsp_autoimport_completions(config, workspace, doc, com_position)
-#     assert len(completions) == 0
-#     workspace.rm_document(DOC_URI)
-
-
 class TestShouldInsert:
     def test_dot(self):
         assert not should_insert("""str.""", 4)
@@ -233,3 +214,74 @@ class sfa:
     """
     results = get_names(jedi.Script(code=source))
     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")
+    )

test/test_language_server.py

@@ -5,7 +5,7 @@
 import time
 import sys
 
-from test.test_utils import ClientServerPair
+from test.test_utils import ClientServerPair, send_initialize_request
 
 from flaky import flaky
 from pylsp_jsonrpc.exceptions import JsonRpcMethodNotFound
@@ -73,14 +73,7 @@ def test_not_exit_without_check_parent_process_flag(
     client_server_pair,
 ):
     client, _ = client_server_pair
-    response = client._endpoint.request(
-        "initialize",
-        {
-            "processId": 1234,
-            "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
-        },
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    response = send_initialize_request(client)
     assert "capabilities" in response