wip

Author: jasongrout

pylsp/python_lsp.py

@@ -585,10 +585,10 @@ 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):
+    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.
-        workspace = self._match_uri_to_workspace(notebookDocument['uri'])
+        workspace = self._match_uri_to_workspace(cellDocument.notebook_uri)
 
         random_uri = str(uuid.uuid4())
         # cell_list helps us map the diagnostics back to the correct cell later.
@@ -609,6 +609,9 @@ def m_notebook_document__definition(self, cellDocument=None, position=None, **_k
                 '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
@@ -619,11 +622,10 @@ def m_notebook_document__definition(self, cellDocument=None, position=None, **_k
 
         # 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:
-            # TODO: adjust position to temp document
-            # position = 
             definitions = self.definitions(random_uri, position)
+            log.info(f'Got definitions: {definitions}')
 
             # {
             #     'uri': uris.uri_with(document.uri, path=str(d.module_path)),
@@ -632,19 +634,21 @@ def m_notebook_document__definition(self, cellDocument=None, position=None, **_k
             #         'end': {'line': d.line - 1, 'character': d.column + len(d.name)},
             #     }
             # }
+            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
+            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.m_notebook_document__definition(textDocument, position, **_kwargs)
+            return self._cell_document__definition(textDocument, 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": "x=1",
                     },
                     {
                         "uri": "cell_2_uri",
@@ -634,8 +634,6 @@ def test_notebook_definition(
         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",
@@ -649,7 +647,9 @@ def test_notebook_definition(
                 }
             },
         )
-        wait_for_condition(lambda: mock_notify.call_count >= 2)
+        raise ValueError(mock_notify.mock_calls)
+        wait_for_condition(lambda: mock_notify.call_count >= 1)
+        return
 
 
         # TODO: definition return