Add support for notebook document completions (python-lsp#486)

pylsp/python_lsp.py

@@ -687,7 +687,36 @@ def m_text_document__code_action(
     def m_text_document__code_lens(self, textDocument=None, **_kwargs):
         return self.code_lens(textDocument["uri"])
 
+    def _cell_document__completion(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"]
+
+            completions = self.completions(temp_uri, position)
+
+            # Translate temp_uri locations to cell document locations
+            for item in completions.get("items", []):
+                if item.get("data", {}).get("doc_uri") == temp_uri:
+                    item["data"]["doc_uri"] = cellDocument.uri
+
+            return completions
+
     def m_text_document__completion(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__completion(document, position, **_kwargs)
         return self.completions(textDocument["uri"], position)
 
     def _cell_document__definition(self, cellDocument, position=None, **_kwargs):

test/test_notebook_document.py

@@ -488,3 +488,46 @@ def test_notebook_definition(client_server_pair):
             },
         }
     ]
+
+
+@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+def test_notebook_completion(client_server_pair):
+    """
+    Tests that completions work across cell boundaries for notebook document support
+    """
+    client, server = client_server_pair
+    send_initialize_request(client)
+
+    # Open notebook
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        send_notebook_did_open(
+            client, ["answer_to_life_universe_everything = 42", "answer_"]
+        )
+        # 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/completion",
+        {
+            "textDocument": {
+                "uri": "cell_2_uri",
+            },
+            "position": {"line": 0, "character": 7},
+        },
+    )
+    result = future.result(CALL_TIMEOUT_IN_SECONDS)
+    assert result == {
+        "isIncomplete": False,
+        "items": [
+            {
+                "data": {"doc_uri": "cell_2_uri"},
+                "insertText": "answer_to_life_universe_everything",
+                "kind": 6,
+                "label": "answer_to_life_universe_everything",
+                "sortText": "aanswer_to_life_universe_everything",
+            },
+        ],
+    }