diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py index faa93423..63d5478d 100644 --- a/pylsp/python_lsp.py +++ b/pylsp/python_lsp.py @@ -16,7 +16,7 @@ from . import lsp, _utils, uris from .config import config -from .workspace import Workspace, Document, Notebook +from .workspace import Workspace, Document, Notebook, Cell from ._version import __version__ log = logging.getLogger(__name__) @@ -541,6 +541,7 @@ def m_notebook_document__did_open( for cell in cellTextDocuments or []: workspace.put_cell_document( cell["uri"], + notebookDocument["uri"], cell["languageId"], cell["text"], version=cell.get("version"), @@ -593,6 +594,7 @@ def m_notebook_document__did_change( for cell_document in structure["didOpen"]: workspace.put_cell_document( cell_document["uri"], + notebookDocument["uri"], cell_document["languageId"], cell_document["text"], cell_document.get("version"), @@ -671,7 +673,46 @@ def m_text_document__code_lens(self, textDocument=None, **_kwargs): def m_text_document__completion(self, textDocument=None, position=None, **_kwargs): 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) + 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"] + + definitions = self.definitions(temp_uri, position) + + # Translate temp_uri locations to cell document locations + for definition in definitions: + if definition["uri"] == temp_uri: + # Find the cell the start line is in and adjust the uri and line numbers + for cell_uri, data in cell_data.items(): + if ( + data["line_start"] + <= definition["range"]["start"]["line"] + <= data["line_end"] + ): + definition["uri"] = cell_uri + definition["range"]["start"]["line"] -= data["line_start"] + definition["range"]["end"]["line"] -= data["line_start"] + break + + return definitions + def m_text_document__definition(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__definition(document, position, **_kwargs) return self.definitions(textDocument["uri"], position) def m_text_document__document_highlight( diff --git a/pylsp/workspace.py b/pylsp/workspace.py index d9868846..27af5f83 100644 --- a/pylsp/workspace.py +++ b/pylsp/workspace.py @@ -130,6 +130,17 @@ def put_notebook_document( doc_uri, notebook_type, cells, version, metadata ) + @contextmanager + def temp_document(self, source, path=None): + if path is None: + path = self.root_path + uri = uris.from_fs_path(os.path.join(path, str(uuid.uuid4()))) + try: + self.put_document(uri, source) + yield uri + finally: + self.rm_document(uri) + def add_notebook_cells(self, doc_uri, cells, start): self._docs[doc_uri].add_cells(cells, start) @@ -139,9 +150,11 @@ def remove_notebook_cells(self, doc_uri, start, delete_count): def update_notebook_metadata(self, doc_uri, metadata): self._docs[doc_uri].metadata = metadata - def put_cell_document(self, doc_uri, language_id, source, version=None): + def put_cell_document( + self, doc_uri, notebook_uri, language_id, source, version=None + ): self._docs[doc_uri] = self._create_cell_document( - doc_uri, language_id, source, version + doc_uri, notebook_uri, language_id, source, version ) def rm_document(self, doc_uri): @@ -340,11 +353,14 @@ def _create_notebook_document( metadata=metadata, ) - def _create_cell_document(self, doc_uri, language_id, source=None, version=None): + def _create_cell_document( + self, doc_uri, notebook_uri, language_id, source=None, version=None + ): # TODO: remove what is unnecessary here. path = uris.to_fs_path(doc_uri) return Cell( doc_uri, + notebook_uri=notebook_uri, language_id=language_id, workspace=self, source=source, @@ -585,6 +601,26 @@ def add_cells(self, new_cells: List, start: int) -> None: def remove_cells(self, start: int, delete_count: int) -> None: del self.cells[start : start + delete_count] + def cell_data(self): + """Extract current cell data. + + Returns a dict (ordered by cell position) where the key is the cell uri and the + value is a dict with line_start, line_end, and source attributes. + """ + cell_data = {} + offset = 0 + for cell in self.cells: + cell_uri = cell["document"] + cell_document = self.workspace.get_cell_document(cell_uri) + num_lines = cell_document.line_count + cell_data[cell_uri] = { + "line_start": offset, + "line_end": offset + num_lines - 1, + "source": cell_document.source, + } + offset += num_lines + return cell_data + class Cell(Document): """ @@ -599,6 +635,7 @@ class Cell(Document): def __init__( self, uri, + notebook_uri, language_id, workspace, source=None, @@ -611,6 +648,7 @@ def __init__( uri, workspace, source, version, local, extra_sys_path, rope_project_builder ) self.language_id = language_id + self.notebook_uri = notebook_uri @property @lock diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py index d632a660..e8e7ac75 100644 --- a/test/test_notebook_document.py +++ b/test/test_notebook_document.py @@ -544,3 +544,75 @@ def test_notebook__did_close( ) wait_for_condition(lambda: mock_notify.call_count >= 2) assert len(server.workspace.documents) == 0 + + +@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) + + # 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", + }, + ], + }, + ) + # 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/definition", + { + "textDocument": { + "uri": "cell_2_uri", + }, + "position": {"line": 0, "character": 1}, + }, + ) + result = future.result(CALL_TIMEOUT_IN_SECONDS) + assert result == [ + { + "uri": "cell_1_uri", + "range": { + "start": {"line": 1, "character": 0}, + "end": {"line": 1, "character": 1}, + }, + } + ] diff --git a/test/test_workspace.py b/test/test_workspace.py index 363ee976..02ef1d22 100644 --- a/test/test_workspace.py +++ b/test/test_workspace.py @@ -7,6 +7,7 @@ DOC_URI = uris.from_fs_path(__file__) +NOTEBOOK_URI = uris.from_fs_path("notebook_uri") def path_as_uri(path): @@ -29,7 +30,7 @@ def test_put_notebook_document(pylsp): def test_put_cell_document(pylsp): - pylsp.workspace.put_cell_document(DOC_URI, "python", "content") + pylsp.workspace.put_cell_document(DOC_URI, NOTEBOOK_URI, "python", "content") assert DOC_URI in pylsp.workspace._docs