diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py index 52a22a3e..a31e7612 100644 --- a/pylsp/python_lsp.py +++ b/pylsp/python_lsp.py @@ -687,7 +687,36 @@ def m_text_document__code_action( def m_text_document__code_lens(self, textDocument=None, **_kwargs): return self.code_lens(textDocument["uri"]) + 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 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__definition(self, cellDocument, position=None, **_kwargs): diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py index c63d2791..f97088b3 100644 --- a/test/test_notebook_document.py +++ b/test/test_notebook_document.py @@ -488,3 +488,46 @@ 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", + }, + ], + }