Fix go to definition for notebook documents

Fix test

Author: jasongrout

pylsp/python_lsp.py

@@ -586,9 +586,12 @@ def m_text_document__completion(self, textDocument=None, position=None, **_kwarg
         return self.completions(textDocument['uri'], position)
 
     def _cell_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.
+        # 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.
@@ -636,19 +639,26 @@ def _cell_document__definition(self, cellDocument=None, position=None, **_kwargs
             # }
             print(definitions)
             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
+                # 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):
-        if isinstance(textDocument, Cell):
-            # actually, test to see if the document is a cell document
-            return self._cell_document__definition(textDocument, position, **_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):

test/test_notebook_document.py

@@ -619,7 +619,7 @@ def test_notebook_definition(
                     {
                         "uri": "cell_1_uri",
                         "languageId": "python",
-                        "text": "x=1",
+                        "text": "y=2\nx=1",
                     },
                     {
                         "uri": "cell_2_uri",
@@ -629,52 +629,35 @@ def test_notebook_definition(
                 ],
             },
         )
+        # 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
 
-    with patch.object(server._endpoint, "notify") as mock_notify:
-        client._endpoint.notify(
-            "textDocument/definition",
-            {
-                "textDocument": {
-                    "uri": "cell_2_uri",
-                },
-                "position": {
-                    "line": 0,
-                    "character": 1
-                }
+    future = client._endpoint.request(
+        "textDocument/definition",
+        {
+            "textDocument": {
+                "uri": "cell_2_uri",
             },
-        )
-        raise ValueError(mock_notify.mock_calls)
-        wait_for_condition(lambda: mock_notify.call_count >= 1)
-        return
-
-
-        # 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)
+            "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
+            }
+        }
+    }]