Notebook protocol go-to-definition support

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(

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

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},
+            },
+        }
+    ]

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