rope autoimport - ignore names from the notebook document when doing …

…completions on a cell

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,9 +168,13 @@ 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)
+    )
+    log.debug("autoimport: ignored names: %s", ignored_names)
     autoimport = workspace._rope_autoimport(rope_config)
     suggestions = list(autoimport.search_full(word, ignored_names=ignored_names))
+    log.debug("autoimport: suggestions: %s", suggestions)
     results = list(
         sorted(
             _process_statements(suggestions, document.uri, word, autoimport, document),

pylsp/python_lsp.py

@@ -394,7 +394,18 @@ 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):
+            log.debug("Getting ignored names from notebook document")
+            # 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()
+            log.debug("Got ignored names from notebook document: %s", ignored_names)
+        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,15 @@ def cell_data(self):
             offset += num_lines
         return cell_data
 
+    @lock
+    def jedi_names(self, all_scopes=False, definitions=True, references=False):
+        """Get all names to ignore from all cells."""
+        names = set()
+        for cell in self.cells:
+            cell_document = self.workspace.get_cell_document(cell["document"])
+            names.update(cell_document.jedi_names(all_scopes, definitions, references))
+        return set(name.name for name in names)
+
 
 class Cell(Document):
     """

test/fixtures.py

@@ -5,7 +5,8 @@
 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 pylsp_jsonrpc
 
 import pytest
 from pylsp_jsonrpc.dispatchers import MethodDispatcher
@@ -24,7 +25,6 @@
 def main():
     print sys.stdin.read()
 """
-CALL_TIMEOUT_IN_SECONDS = 30
 
 
 class FakeEditorMethodsMixin:
@@ -175,8 +175,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,18 @@
 # Copyright 2022- Python Language Server Contributors.
 
 from typing import Dict, List
-from unittest.mock import Mock
+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, uris
+from pylsp import lsp
+
 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,6 +21,7 @@
 from pylsp.plugins.rope_autoimport import pylsp_initialize
 from pylsp.workspace import Workspace
 
+
 DOC_URI = uris.from_fs_path(__file__)
 
 
@@ -39,7 +45,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 +149,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 +209,40 @@ 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:
+        send_notebook_did_open(client, ["import os", "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
+
+    suggestions = server.completions("cell_1_uri", {"line": 0, "character": 2}).get(
+        "items"
+    )
+    assert not any(
+        suggestion for suggestion in suggestions if suggestion.get("label", "") == "os"
+    )
+
+    suggestions = server.completions("cell_2_uri", {"line": 0, "character": 3}).get(
+        "items"
+    )
+    assert any(
+        suggestion for suggestion in suggestions if suggestion.get("label", "") == "sys"
+    )

test/test_notebook_document.py

@@ -4,7 +4,7 @@
 import time
 
 from unittest.mock import patch, call
-from test.fixtures import CALL_TIMEOUT_IN_SECONDS
+from test.test_utils import CALL_TIMEOUT_IN_SECONDS
 import pytest
 from pylsp.workspace import Notebook
 
@@ -29,7 +29,6 @@ def test_initialize(client_server_pair):
         {
             "processId": 1234,
             "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
         },
     ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
     assert server.workspace is not None
@@ -100,7 +99,6 @@ def test_notebook_document__did_open(
         {
             "processId": 1234,
             "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
         },
     ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
 
@@ -264,7 +262,6 @@ def test_notebook_document__did_change(
         {
             "processId": 1234,
             "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
         },
     ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
 
@@ -536,7 +533,6 @@ def test_notebook__did_close(
         {
             "processId": 1234,
             "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
         },
     ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
 
@@ -608,7 +604,6 @@ def test_notebook_definition(client_server_pair):
         {
             "processId": 1234,
             "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
         },
     ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
 

test/test_utils.py

@@ -6,15 +6,72 @@
 import sys
 from threading import Thread
 import time
+from typing import List
 from unittest import mock
 
 from flaky import flaky
 from docstring_to_markdown import UnknownFormatError
+from build.lib.pylsp.lsp import NotebookCellKind
 
 from pylsp import _utils
 from pylsp.python_lsp import PythonLSPServer, start_io_lang_server
 
 
+CALL_TIMEOUT_IN_SECONDS = 30
+
+
+def send_notebook_did_open(client, cells: List[str]):
+    """
+    Sends a notebookDocument/didOpen notification with the given python cells.
+
+    The notebook has the uri "notebook_uri" and the cells have the uris
+    "cell_0_uri", "cell_1_uri", etc.
+    """
+    client._endpoint.notify(
+        "notebookDocument/didOpen", notebook_with_python_cells(cells)
+    )
+
+
+def notebook_with_python_cells(cells: List[str]):
+    """
+    Create a notebook document with the given python cells.
+
+    The notebook has the uri "notebook_uri" and the cells have the uris
+    "cell_0_uri", "cell_1_uri", etc.
+    """
+    return {
+        "notebookDocument": {
+            "uri": "notebook_uri",
+            "notebookType": "jupyter-notebook",
+            "cells": [
+                {
+                    "kind": NotebookCellKind.Code,
+                    "document": f"cell_{i}_uri",
+                }
+                for i in range(len(cells))
+            ],
+        },
+        "cellTextDocuments": [
+            {
+                "uri": f"cell_{i}_uri",
+                "languageId": "python",
+                "text": cell,
+            }
+            for i, cell in enumerate(cells)
+        ],
+    }
+
+
+def send_initialize_request(client):
+    client._endpoint.request(
+        "initialize",
+        {
+            "processId": 1234,
+            "rootPath": os.path.dirname(__file__),
+        },
+    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+
+
 def start(obj):
     obj.start()