From fc40eb17d5085f8aaecb1c07a78461bd56bf9bb4 Mon Sep 17 00:00:00 2001 From: Stephen Macke Date: Thu, 9 Nov 2023 09:11:25 -0800 Subject: [PATCH 1/4] attempt to support completions for notebook document --- pylsp/python_lsp.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py index 52a22a3e..7eb8b311 100644 --- a/pylsp/python_lsp.py +++ b/pylsp/python_lsp.py @@ -688,8 +688,37 @@ def m_text_document__code_lens(self, textDocument=None, **_kwargs): return self.code_lens(textDocument["uri"]) def m_text_document__completion(self, textDocument=None, position=None, **_kwargs): + # textDocument here is just a dict with a uri + workspace = self._match_uri_to_workspace(textDocument["uri"]) + document = workspace.get_document(textDocument["uri"]) + if isinstance(document, Cell): + return self._cell_document__completion(document, position, **_kwargs) return self.completions(textDocument["uri"], position) + def _cell_document__completion(self, cellDocument, position=None, **_kwargs): + workspace = self._match_uri_to_workspace(cellDocument.notebook_uri) + notebookDocument = workspace.get_maybe_document(cellDocument.notebook_uri) + if notebookDocument is None: + raise ValueError("Invalid notebook document") + + cell_data = notebookDocument.cell_data() + + # Concatenate all cells to be a single temporary document + total_source = "\n".join(data["source"] for data in cell_data.values()) + with workspace.temp_document(total_source) as temp_uri: + # update position to be the position in the temp document + if position is not None: + position["line"] += cell_data[cellDocument.uri]["line_start"] + + completions = self.completions(temp_uri, position) + + # Translate temp_uri locations to cell document locations + for item in completions.get("items", []): + if item.get("data", {}).get("doc_uri") == temp_uri: + item["data"]["doc_uri"] = cellDocument.uri + + return completions + def _cell_document__definition(self, cellDocument, position=None, **_kwargs): workspace = self._match_uri_to_workspace(cellDocument.notebook_uri) notebookDocument = workspace.get_maybe_document(cellDocument.notebook_uri) From 4e6a224ceef94d5a298320b6ed9f9be536021ed9 Mon Sep 17 00:00:00 2001 From: Stephen Macke Date: Sun, 12 Nov 2023 15:27:22 -0800 Subject: [PATCH 2/4] add test --- pylsp/python_lsp.py | 20 +++++++++-------- test/test_notebook_document.py | 41 ++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 9 deletions(-) diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py index 7eb8b311..6d71fcb4 100644 --- a/pylsp/python_lsp.py +++ b/pylsp/python_lsp.py @@ -687,14 +687,6 @@ def m_text_document__code_action( def m_text_document__code_lens(self, textDocument=None, **_kwargs): return self.code_lens(textDocument["uri"]) - def m_text_document__completion(self, textDocument=None, position=None, **_kwargs): - # textDocument here is just a dict with a uri - workspace = self._match_uri_to_workspace(textDocument["uri"]) - document = workspace.get_document(textDocument["uri"]) - if isinstance(document, Cell): - return self._cell_document__completion(document, position, **_kwargs) - return self.completions(textDocument["uri"], position) - def _cell_document__completion(self, cellDocument, position=None, **_kwargs): workspace = self._match_uri_to_workspace(cellDocument.notebook_uri) notebookDocument = workspace.get_maybe_document(cellDocument.notebook_uri) @@ -719,6 +711,15 @@ def _cell_document__completion(self, cellDocument, position=None, **_kwargs): return completions + def m_text_document__completion(self, textDocument=None, position=None, **_kwargs): + # textDocument here is just a dict with a uri + workspace = self._match_uri_to_workspace(textDocument["uri"]) + document = workspace.get_document(textDocument["uri"]) + if isinstance(document, Cell): + return self._cell_document__completion(document, position, **_kwargs) + else: + return self.completions(textDocument["uri"], position) + def _cell_document__definition(self, cellDocument, position=None, **_kwargs): workspace = self._match_uri_to_workspace(cellDocument.notebook_uri) notebookDocument = workspace.get_maybe_document(cellDocument.notebook_uri) @@ -759,7 +760,8 @@ def m_text_document__definition(self, textDocument=None, position=None, **_kwarg document = workspace.get_document(textDocument["uri"]) if isinstance(document, Cell): return self._cell_document__definition(document, position, **_kwargs) - return self.definitions(textDocument["uri"], position) + else: + return self.definitions(textDocument["uri"], position) def m_text_document__document_highlight( self, textDocument=None, position=None, **_kwargs diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py index c63d2791..28969aa2 100644 --- a/test/test_notebook_document.py +++ b/test/test_notebook_document.py @@ -488,3 +488,44 @@ def test_notebook_definition(client_server_pair): }, } ] + + +@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows") +def test_notebook_completion(client_server_pair): + """ + Tests that completions work across cell boundaries for notebook document support + """ + client, server = client_server_pair + send_initialize_request(client) + + # Open notebook + with patch.object(server._endpoint, "notify") as mock_notify: + send_notebook_did_open(client, ["answer_to_life_universe_everything = 42", "answer_"]) + # wait for expected diagnostics messages + 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"]: + assert uri in server.workspace.documents + + future = client._endpoint.request( + "textDocument/completion", + { + "textDocument": { + "uri": "cell_2_uri", + }, + "position": {"line": 0, "character": 7}, + }, + ) + result = future.result(CALL_TIMEOUT_IN_SECONDS) + assert result == { + 'isIncomplete': False, + 'items': [ + { + 'data': {'doc_uri': 'cell_2_uri'}, + 'insertText': 'answer_to_life_universe_everything', + 'kind': 6, + 'label': 'answer_to_life_universe_everything', + 'sortText': 'aanswer_to_life_universe_everything', + }, + ], + } From 4e1d58e1528a98ef5a45e81191e29c1aa6626b39 Mon Sep 17 00:00:00 2001 From: Stephen Macke Date: Sun, 12 Nov 2023 15:28:39 -0800 Subject: [PATCH 3/4] blacken --- test/test_notebook_document.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py index 28969aa2..f97088b3 100644 --- a/test/test_notebook_document.py +++ b/test/test_notebook_document.py @@ -500,7 +500,9 @@ def test_notebook_completion(client_server_pair): # Open notebook with patch.object(server._endpoint, "notify") as mock_notify: - send_notebook_did_open(client, ["answer_to_life_universe_everything = 42", "answer_"]) + send_notebook_did_open( + client, ["answer_to_life_universe_everything = 42", "answer_"] + ) # wait for expected diagnostics messages wait_for_condition(lambda: mock_notify.call_count >= 2) assert len(server.workspace.documents) == 3 @@ -518,14 +520,14 @@ def test_notebook_completion(client_server_pair): ) result = future.result(CALL_TIMEOUT_IN_SECONDS) assert result == { - 'isIncomplete': False, - 'items': [ + "isIncomplete": False, + "items": [ { - 'data': {'doc_uri': 'cell_2_uri'}, - 'insertText': 'answer_to_life_universe_everything', - 'kind': 6, - 'label': 'answer_to_life_universe_everything', - 'sortText': 'aanswer_to_life_universe_everything', + "data": {"doc_uri": "cell_2_uri"}, + "insertText": "answer_to_life_universe_everything", + "kind": 6, + "label": "answer_to_life_universe_everything", + "sortText": "aanswer_to_life_universe_everything", }, ], } From 729a9d9618d8a88c87d58b6dcd360e02f3dbb1e5 Mon Sep 17 00:00:00 2001 From: Stephen Macke Date: Sun, 12 Nov 2023 15:35:49 -0800 Subject: [PATCH 4/4] fix lints --- pylsp/python_lsp.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py index 6d71fcb4..a31e7612 100644 --- a/pylsp/python_lsp.py +++ b/pylsp/python_lsp.py @@ -717,8 +717,7 @@ def m_text_document__completion(self, textDocument=None, position=None, **_kwarg document = workspace.get_document(textDocument["uri"]) if isinstance(document, Cell): return self._cell_document__completion(document, position, **_kwargs) - else: - return self.completions(textDocument["uri"], position) + return self.completions(textDocument["uri"], position) def _cell_document__definition(self, cellDocument, position=None, **_kwargs): workspace = self._match_uri_to_workspace(cellDocument.notebook_uri) @@ -760,8 +759,7 @@ def m_text_document__definition(self, textDocument=None, position=None, **_kwarg document = workspace.get_document(textDocument["uri"]) if isinstance(document, Cell): return self._cell_document__definition(document, position, **_kwargs) - else: - return self.definitions(textDocument["uri"], position) + return self.definitions(textDocument["uri"], position) def m_text_document__document_highlight( self, textDocument=None, position=None, **_kwargs