Add go to definition support for notebook cell documents

Author: jasongrout

README.md

@@ -165,8 +165,8 @@ Dev install
 
 ```
 # create conda env
-cc python-lsp-server
-ca python-lsp-server
+conda create python-lsp-server
+conda activate python-lsp-server
 
 pip install ".[all]"
 pip install ".[websockets]"

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__)
@@ -480,7 +480,7 @@ def m_notebook_document__did_open(self, notebookDocument=None, cellTextDocuments
                                         cells=notebookDocument['cells'], version=notebookDocument.get('version'),
                                         metadata=notebookDocument.get('metadata'))
         for cell in (cellTextDocuments or []):
-            workspace.put_cell_document(cell['uri'], cell['languageId'], cell['text'], version=cell.get('version'))
+            workspace.put_cell_document(cell['uri'], notebookDocument['uri'], cell['languageId'], cell['text'], version=cell.get('version'))
         self.lint(notebookDocument['uri'], is_saved=True)
 
     def m_notebook_document__did_close(self, notebookDocument=None, cellTextDocuments=None, **_kwargs):
@@ -521,7 +521,7 @@ def m_notebook_document__did_change(self, notebookDocument=None, change=None, **
                     # Case 2
                     # Cell documents
                     for cell_document in structure['didOpen']:
-                        workspace.put_cell_document(cell_document['uri'], cell_document['languageId'],
+                        workspace.put_cell_document(cell_document['uri'], notebookDocument['uri'], cell_document['languageId'],
                                                     cell_document['text'], cell_document.get('version'))
                     # Cell metadata which is added to Notebook
                     workspace.add_notebook_cells(notebookDocument['uri'], notebook_cell_array_change['cells'], start)
@@ -585,7 +585,80 @@ 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=None, position=None, **_kwargs):
+        # First, we create a temp TextDocument to send to the hook that represents the whole notebook
+        # contents.
+        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")
+
+        random_uri = str(uuid.uuid4())
+        # cell_list helps us map the diagnostics back to the correct cell later.
+        cell_list: List[Dict[str, Any]] = []
+
+        offset = 0
+        total_source = ""
+        for cell in notebookDocument.cells:
+            cell_uri = cell['document']
+            cell_document = workspace.get_cell_document(cell_uri)
+
+            num_lines = cell_document.line_count
+
+            data = {
+                'uri': cell_uri,
+                'line_start': offset,
+                'line_end': offset + num_lines - 1,
+                'source': cell_document.source
+            }
+
+            if position is not None and cell_uri == cellDocument.uri:
+                position['line'] += offset
+
+            cell_list.append(data)
+            if offset == 0:
+                total_source = cell_document.source
+            else:
+                total_source += ("\n" + cell_document.source)
+
+            offset += num_lines
+
+        # TODO: make a workspace temp document context manager that yields the random uri and cleans up afterwards
+        workspace.put_document(random_uri, total_source)
+        log.info(f'Making new document {random_uri}')
+        try:
+            definitions = self.definitions(random_uri, position)
+            log.info(f'Got definitions: {definitions}')
+
+            # {
+            #     'uri': uris.uri_with(document.uri, path=str(d.module_path)),
+            #     'range': {
+            #         'start': {'line': d.line - 1, 'character': d.column},
+            #         'end': {'line': d.line - 1, 'character': d.column + len(d.name)},
+            #     }
+            # }
+            print(definitions)
+            for definition in definitions:
+                # TODO: a better test for if a definition is the random_uri
+                if random_uri in definition['uri']:
+                    # Find the cell the start is in
+                    for cell in cell_list:
+                        # TODO: perhaps it is more correct to check definition['range']['end']['line'] <= cell['line_end'], but 
+                        # that would mess up if a definition was split over cells
+                        if cell['line_start'] <= definition['range']['start']['line'] <= cell['line_end']:
+                            definition['uri'] = cell['uri']
+                            definition['range']['start']['line'] -= cell['line_start']
+                            definition['range']['end']['line'] -= cell['line_start']
+            return definitions
+        finally:
+            workspace.rm_document(random_uri)
+
     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(self, textDocument=None, position=None, **_kwargs):

pylsp/workspace.py

@@ -128,8 +128,8 @@ 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):
-        self._docs[doc_uri] = self._create_cell_document(doc_uri, language_id, source, version)
+    def put_cell_document(self, doc_uri, notebook_uri, language_id, source, version=None):
+        self._docs[doc_uri] = self._create_cell_document(doc_uri, notebook_uri, language_id, source, version)
 
     def rm_document(self, doc_uri):
         self._docs.pop(doc_uri)
@@ -288,11 +288,12 @@ def _create_notebook_document(self, doc_uri, notebook_type, cells, version=None,
             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,
@@ -511,10 +512,11 @@ def remove_cells(self, start: int, delete_count: int) -> None:
 class Cell(Document):
     """Represents a cell in a notebook."""
 
-    def __init__(self, uri, language_id, workspace, source=None, version=None, local=True, extra_sys_path=None,
+    def __init__(self, uri, notebook_uri, language_id, workspace, source=None, version=None, local=True, extra_sys_path=None,
                  rope_project_builder=None):
         super().__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

@@ -582,3 +582,82 @@ def test_notebook__did_close(
         )
         wait_for_condition(lambda: mock_notify.call_count >= 2)
         assert len(server.workspace.documents) == 0
+
+def test_notebook_definition(
+    client_server_pair,
+):  # pylint: disable=redefined-outer-name
+    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
+            }
+        }
+    }]