WIP 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,66 @@ 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 m_notebook_document__definition(self, cellDocument=None, position=None, **_kwargs):
+        # First, we create a temp TextDocument that represents the whole notebook
+        # contents. We'll use this to send to the hook.
+        workspace = self._match_uri_to_workspace(notebookDocument['uri'])
+
+        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
+            }
+
+            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)
+
+        try:
+            # TODO: adjust position to temp document
+            # position = 
+            definitions = self.definitions(random_uri, position)
+
+            # {
+            #     '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)},
+            #     }
+            # }
+            for definition in definitions:
+                if definition['uri'] == random_uri:
+                    # find what cell the start is in
+                    # make sure the end is inside the cell's line_end
+                    # subtract that cell's line_start from both definition start and end
+                    pass
+        finally:
+            workspace.rm_document(random_uri)
+
     def m_text_document__definition(self, textDocument=None, position=None, **_kwargs):
+        if isinstance(textDocument, Cell):
+            # actually, test to see if the document is a cell document
+            return self.m_notebook_document__definition(textDocument, 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

pyproject.toml

@@ -56,7 +56,7 @@ test = [
     "numpy",
     "pandas",
     "matplotlib",
-    "pyqt5",
+#    "pyqt5",
     "flaky",
 ]
 

test/test_notebook_document.py

@@ -582,3 +582,99 @@ 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": "x = 1",
+                    },
+                    {
+                        "uri": "cell_2_uri",
+                        "languageId": "python",
+                        "text": "x",
+                    },
+                ],
+            },
+        )
+        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
+
+
+    # TODO: send definition request
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        client._endpoint.notify(
+            "textDocument/definition",
+            {
+                "textDocument": {
+                    "uri": "cell_2_uri",
+                },
+                "position": {
+                    "line": 0,
+                    "character": 1
+                }
+            },
+        )
+        wait_for_condition(lambda: mock_notify.call_count >= 2)
+
+
+        # TODO: definition return
+        expected_call_args = [
+            call(
+                "textDocument/publishDiagnostics",
+                params={"uri": "cell_1_uri", "diagnostics": []},
+            ),
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_3_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pycodestyle",
+                            "range": {
+                                "start": {"line": 0, "character": 8},
+                                "end": {"line": 0, "character": 8},
+                            },
+                            "message": "W292 no newline at end of file",
+                            "code": "W292",
+                            "severity": 2,
+                        }
+                    ],
+                },
+            ),
+        ]
+        mock_notify.assert_has_calls(expected_call_args)