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..6997b594 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,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), diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py index 4c3ea0d2..222239a0 100644 --- a/pylsp/python_lsp.py +++ b/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): diff --git a/pylsp/workspace.py b/pylsp/workspace.py index a08e09c6..09559004 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,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): """ diff --git a/test/fixtures.py b/test/fixtures.py index 03d0f824..b1485eac 100644 --- a/test/fixtures.py +++ b/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") diff --git a/test/plugins/test_autoimport.py b/test/plugins/test_autoimport.py index dbb6f7a4..43feccdb 100644 --- a/test/plugins/test_autoimport.py +++ b/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" + ) diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py index 15f187d3..a268ea55 100644 --- a/test/test_notebook_document.py +++ b/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) diff --git a/test/test_utils.py b/test/test_utils.py index bc2782dc..98451182 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 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()