Copy LAST_JEDI_COMPLETIONS to cell document so that completionItem/resolve will work (#663)

hjr265 · web-flow · commit 5c8cfacbcd1f · 2025-08-25T11:48:47.000-05:00
diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
@@ -722,6 +722,12 @@ def _cell_document__completion(self, cellDocument, position=None, **_kwargs):
                 if item.get("data", {}).get("doc_uri") == temp_uri:
                     item["data"]["doc_uri"] = cellDocument.uri
 
+            # Copy LAST_JEDI_COMPLETIONS to cell document so that completionItem/resolve will work
+            tempDocument = workspace.get_document(temp_uri)
+            cellDocument.shared_data["LAST_JEDI_COMPLETIONS"] = (
+                tempDocument.shared_data.get("LAST_JEDI_COMPLETIONS", None)
+            )
+
             return completions
 
     def m_text_document__completion(self, textDocument=None, position=None, **_kwargs):
diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py
@@ -530,3 +530,71 @@ def test_notebook_completion(client_server_pair) -> None:
             },
         ],
     }
+
+
+@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+def test_notebook_completion_resolve(client_server_pair) -> None:
+    """
+    Tests that completion item resolve works correctly
+    """
+    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,
+            [
+                "def answer():\n\t'''Returns an important number.'''\n\treturn 42",
+                "ans",
+            ],
+        )
+        # 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": 3},
+        },
+    )
+    result = future.result(CALL_TIMEOUT_IN_SECONDS)
+    assert result == {
+        "isIncomplete": False,
+        "items": [
+            {
+                "data": {"doc_uri": "cell_2_uri"},
+                "insertText": "answer",
+                "kind": 3,
+                "label": "answer()",
+                "sortText": "aanswer",
+            },
+        ],
+    }
+
+    future = client._endpoint.request(
+        "completionItem/resolve",
+        {
+            "data": {"doc_uri": "cell_2_uri"},
+            "label": "answer()",
+        },
+    )
+    result = future.result(CALL_TIMEOUT_IN_SECONDS)
+    del result["detail"]  # The value of this is unpredictable.
+    assert result == {
+        "data": {"doc_uri": "cell_2_uri"},
+        "insertText": "answer",
+        "kind": 3,
+        "label": "answer()",
+        "sortText": "aanswer",
+        "documentation": {
+            "kind": "markdown",
+            "value": "```python\nanswer()\n```\n\n\nReturns an important number.",
+        },
+    }
