From 7de626f8e00fb619583dee2ae1d2c1ca3e8b2c21 Mon Sep 17 00:00:00 2001 From: tkrabel-db <91616041+tkrabel-db@users.noreply.github.com> Date: Thu, 19 Oct 2023 02:24:24 +0200 Subject: [PATCH] Ignore notebook names on cell completion for autoimport (#466) --- pylsp/hookspecs.py | 2 +- pylsp/plugins/rope_autoimport.py | 12 +- pylsp/python_lsp.py | 11 +- pylsp/workspace.py | 26 ++++ test/fixtures.py | 20 ++- test/plugins/test_autoimport.py | 130 ++++++++++++------ test/test_language_server.py | 11 +- test/test_notebook_document.py | 220 +++---------------------------- test/test_utils.py | 57 ++++++++ 9 files changed, 228 insertions(+), 261 deletions(-) diff --git a/pylsp/hookspecs.py b/pylsp/hookspecs.py index d732b1d0..9c9cf387 100644 --- a/pylsp/hookspecs.py +++ b/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 diff --git a/pylsp/plugins/rope_autoimport.py b/pylsp/plugins/rope_autoimport.py index 6c4784aa..1caab35d 100644 --- a/pylsp/plugins/rope_autoimport.py +++ b/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( diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py index 4c3ea0d2..760ad974 100644 --- a/pylsp/python_lsp.py +++ b/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): diff --git a/pylsp/workspace.py b/pylsp/workspace.py index e0e16ef9..fb524c71 100644 --- a/pylsp/workspace.py +++ b/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): """ diff --git a/test/fixtures.py b/test/fixtures.py index 03d0f824..ed6206af 100644 --- a/test/fixtures.py +++ b/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") diff --git a/test/plugins/test_autoimport.py b/test/plugins/test_autoimport.py index dbb6f7a4..b1c46775 100644 --- a/test/plugins/test_autoimport.py +++ b/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") + ) diff --git a/test/test_language_server.py b/test/test_language_server.py index 280a62fa..401b1ceb 100644 --- a/test/test_language_server.py +++ b/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 diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py index 15f187d3..c63d2791 100644 --- a/test/test_notebook_document.py +++ b/test/test_notebook_document.py @@ -1,10 +1,14 @@ # Copyright 2021- Python Language Server Contributors. -import os 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, + send_initialize_request, + send_notebook_did_open, +) + import pytest from pylsp.workspace import Notebook @@ -24,14 +28,7 @@ def wait_for_condition(condition, timeout=CALL_TIMEOUT_IN_SECONDS): @pytest.mark.skipif(IS_WIN, reason="Flaky on Windows") def test_initialize(client_server_pair): client, server = 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 server.workspace is not None selector = response["capabilities"]["notebookDocumentSync"]["notebookSelector"] assert isinstance(selector, list) @@ -41,13 +38,7 @@ def test_initialize(client_server_pair): def test_workspace_did_change_configuration(client_server_pair): """Test that we can update a workspace config w/o error when a notebook is open.""" client, server = client_server_pair - client._endpoint.request( - "initialize", - { - "processId": 1234, - "rootPath": os.path.dirname(__file__), - }, - ).result(timeout=CALL_TIMEOUT_IN_SECONDS) + send_initialize_request(client) assert server.workspace is not None with patch.object(server._endpoint, "notify") as mock_notify: @@ -95,74 +86,12 @@ def test_notebook_document__did_open( client_server_pair, ): client, server = client_server_pair - client._endpoint.request( - "initialize", - { - "processId": 1234, - "rootPath": os.path.dirname(__file__), - "initializationOptions": {}, - }, - ).result(timeout=CALL_TIMEOUT_IN_SECONDS) + send_initialize_request(client) with patch.object(server._endpoint, "notify") as mock_notify: - client._endpoint.notify( - "notebookDocument/didOpen", - { - "notebookDocument": { - "uri": "notebook_uri", - "notebookType": "jupyter-notebook", - "cells": [ - { - "kind": NotebookCellKind.Code, - "document": "cell_1_uri", - }, - { - "kind": NotebookCellKind.Code, - "document": "cell_2_uri", - }, - { - "kind": NotebookCellKind.Code, - "document": "cell_3_uri", - }, - { - "kind": NotebookCellKind.Code, - "document": "cell_4_uri", - }, - { - "kind": NotebookCellKind.Code, - "document": "cell_5_uri", - }, - ], - }, - # Test as many edge cases as possible for the diagnostics message - "cellTextDocuments": [ - { - "uri": "cell_1_uri", - "languageId": "python", - "text": "", - }, - { - "uri": "cell_2_uri", - "languageId": "python", - "text": "\n", - }, - { - "uri": "cell_3_uri", - "languageId": "python", - "text": "\nimport sys\n\nabc\n\n", - }, - { - "uri": "cell_4_uri", - "languageId": "python", - "text": "x", - }, - { - "uri": "cell_5_uri", - "languageId": "python", - "text": "y\n", - }, - ], - }, + # Test as many edge cases as possible for the diagnostics messages + send_notebook_did_open( + client, ["", "\n", "\nimport sys\n\nabc\n\n", "x", "y\n"] ) wait_for_condition(lambda: mock_notify.call_count >= 5) expected_call_args = [ @@ -259,48 +188,11 @@ def test_notebook_document__did_change( client_server_pair, ): client, server = client_server_pair - client._endpoint.request( - "initialize", - { - "processId": 1234, - "rootPath": os.path.dirname(__file__), - "initializationOptions": {}, - }, - ).result(timeout=CALL_TIMEOUT_IN_SECONDS) + send_initialize_request(client) # Open notebook with patch.object(server._endpoint, "notify") as mock_notify: - client._endpoint.notify( - "notebookDocument/didOpen", - { - "notebookDocument": { - "uri": "notebook_uri", - "notebookType": "jupyter-notebook", - "cells": [ - { - "kind": NotebookCellKind.Code, - "document": "cell_1_uri", - }, - { - "kind": NotebookCellKind.Code, - "document": "cell_2_uri", - }, - ], - }, - "cellTextDocuments": [ - { - "uri": "cell_1_uri", - "languageId": "python", - "text": "import sys", - }, - { - "uri": "cell_2_uri", - "languageId": "python", - "text": "", - }, - ], - }, - ) + send_notebook_did_open(client, ["import sys", ""]) wait_for_condition(lambda: mock_notify.call_count >= 2) assert len(server.workspace.documents) == 3 for uri in ["cell_1_uri", "cell_2_uri", "notebook_uri"]: @@ -531,48 +423,11 @@ def test_notebook__did_close( client_server_pair, ): client, server = client_server_pair - client._endpoint.request( - "initialize", - { - "processId": 1234, - "rootPath": os.path.dirname(__file__), - "initializationOptions": {}, - }, - ).result(timeout=CALL_TIMEOUT_IN_SECONDS) + send_initialize_request(client) # Open notebook with patch.object(server._endpoint, "notify") as mock_notify: - client._endpoint.notify( - "notebookDocument/didOpen", - { - "notebookDocument": { - "uri": "notebook_uri", - "notebookType": "jupyter-notebook", - "cells": [ - { - "kind": NotebookCellKind.Code, - "document": "cell_1_uri", - }, - { - "kind": NotebookCellKind.Code, - "document": "cell_2_uri", - }, - ], - }, - "cellTextDocuments": [ - { - "uri": "cell_1_uri", - "languageId": "python", - "text": "import sys", - }, - { - "uri": "cell_2_uri", - "languageId": "python", - "text": "", - }, - ], - }, - ) + send_notebook_did_open(client, ["import sys", ""]) wait_for_condition(lambda: mock_notify.call_count >= 2) assert len(server.workspace.documents) == 3 for uri in ["cell_1_uri", "cell_2_uri", "notebook_uri"]: @@ -603,48 +458,11 @@ def test_notebook__did_close( @pytest.mark.skipif(IS_WIN, reason="Flaky on Windows") def test_notebook_definition(client_server_pair): client, server = client_server_pair - client._endpoint.request( - "initialize", - { - "processId": 1234, - "rootPath": os.path.dirname(__file__), - "initializationOptions": {}, - }, - ).result(timeout=CALL_TIMEOUT_IN_SECONDS) + send_initialize_request(client) # Open notebook with patch.object(server._endpoint, "notify") as mock_notify: - client._endpoint.notify( - "notebookDocument/didOpen", - { - "notebookDocument": { - "uri": "notebook_uri", - "notebookType": "jupyter-notebook", - "cells": [ - { - "kind": NotebookCellKind.Code, - "document": "cell_1_uri", - }, - { - "kind": NotebookCellKind.Code, - "document": "cell_2_uri", - }, - ], - }, - "cellTextDocuments": [ - { - "uri": "cell_1_uri", - "languageId": "python", - "text": "y=2\nx=1", - }, - { - "uri": "cell_2_uri", - "languageId": "python", - "text": "x", - }, - ], - }, - ) + send_notebook_did_open(client, ["y=2\nx=1", "x"]) # wait for expected diagnostics messages wait_for_condition(lambda: mock_notify.call_count >= 2) assert len(server.workspace.documents) == 3 diff --git a/test/test_utils.py b/test/test_utils.py index bc2782dc..fb4a8b81 100644 --- a/test/test_utils.py +++ b/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 pylsp import _utils +from pylsp.lsp import NotebookCellKind 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_1_uri", "cell_2_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_1_uri", "cell_2_uri", etc. + """ + return { + "notebookDocument": { + "uri": "notebook_uri", + "notebookType": "jupyter-notebook", + "cells": [ + { + "kind": NotebookCellKind.Code, + "document": f"cell_{i+1}_uri", + } + for i in range(len(cells)) + ], + }, + "cellTextDocuments": [ + { + "uri": f"cell_{i+1}_uri", + "languageId": "python", + "text": cell, + } + for i, cell in enumerate(cells) + ], + } + + +def send_initialize_request(client): + return client._endpoint.request( + "initialize", + { + "processId": 1234, + "rootPath": os.path.dirname(__file__), + }, + ).result(timeout=CALL_TIMEOUT_IN_SECONDS) + + def start(obj): obj.start()