Repository: python-lsp/python-lsp-server
Pull request #486 (2023-11-12 23:30:56)

Add support for notebook document completions
In this PR, we add support for notebook document completions, so that completions can work across cell boundaries. The implementation is straightforward and closely follows that for definitions.

Test plan: added a unit test that fails without the changes in this PR and passes with the changes.


Patch:
diff --git a/pylsp/plugins/jedi_completion.py b/pylsp/plugins/jedi_completion.py
index 2ecf0be..65f7523 100644
--- a/pylsp/plugins/jedi_completion.py
+++ b/pylsp/plugins/jedi_completion.py
@@ -88,6 +88,7 @@ def pylsp_completions(config, document, position):
             include_params=include_params if c.type in ["class", "function"] else False,
             resolve=resolve_eagerly,
             resolve_label_or_snippet=(i < max_to_resolve),
+            snippet_support=snippet_support,
         )
         for i, c in enumerate(completions)
     ]
@@ -102,6 +103,7 @@ def pylsp_completions(config, document, position):
                     include_params=False,
                     resolve=resolve_eagerly,
                     resolve_label_or_snippet=(i < max_to_resolve),
+                    snippet_support=snippet_support,
                 )
                 completion_dict["kind"] = lsp.CompletionItemKind.TypeParameter
                 completion_dict["label"] += " object"
@@ -116,6 +118,7 @@ def pylsp_completions(config, document, position):
                     include_params=False,
                     resolve=resolve_eagerly,
                     resolve_label_or_snippet=(i < max_to_resolve),
+                    snippet_support=snippet_support,
                 )
                 completion_dict["kind"] = lsp.CompletionItemKind.TypeParameter
                 completion_dict["label"] += " object"
@@ -226,6 +229,7 @@ def _format_completion(
     include_params=True,
     resolve=False,
     resolve_label_or_snippet=False,
+    snippet_support=False,
 ):
     completion = {
         "label": _label(d, resolve_label_or_snippet),
@@ -240,16 +244,20 @@ def _format_completion(
     # Adjustments for file completions
     if d.type == "path":
         path = os.path.normpath(d.name)
-        path = path.replace("\\", "\\\\")
-        path = path.replace("/", "\\/")
 
-        # If the completion ends with os.sep, it means it's a directory. So we add an escaped os.sep
-        # at the end to ease additional file completions.
+        # If the completion ends with os.sep, it means it's a directory. So we add os.sep at the end
+        # to ease additional file completions.
         if d.name.endswith(os.sep):
             if os.name == "nt":
-                path = path + "\\\\"
+                path = path + "\\"
             else:
-                path = path + "\\/"
+                path = path + "/"
+
+        # Escape to prevent conflicts with the code snippets grammer
+        # See also https://github.com/python-lsp/python-lsp-server/issues/373
+        if snippet_support:
+            path = path.replace("\\", "\\\\")
+            path = path.replace("/", "\\/")
 
         completion["insertText"] = path
 
diff --git a/pylsp/plugins/pylint_lint.py b/pylsp/plugins/pylint_lint.py
index 67f68fa..4dbb137 100644
--- a/pylsp/plugins/pylint_lint.py
+++ b/pylsp/plugins/pylint_lint.py
@@ -286,7 +286,7 @@ def _run_pylint_stdio(pylint_executable, document, flags):
         p = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
     except IOError:
         log.debug("Can't execute %s. Trying with 'python -m pylint'", pylint_executable)
-        cmd = ["python", "-m", "pylint"]
+        cmd = [sys.executable, "-m", "pylint"]
         cmd.extend(flags)
         cmd.extend(["--from-stdin", document.path])
         p = Popen(  # pylint: disable=consider-using-with
diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
index 52a22a3..a31e761 100644
--- a/pylsp/python_lsp.py
+++ b/pylsp/python_lsp.py
@@ -687,7 +687,36 @@ class PythonLSPServer(MethodDispatcher):
     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 patch:
diff --git a/test/plugins/test_completion.py b/test/plugins/test_completion.py
index 00a54eb..f64e81e 100644
--- a/test/plugins/test_completion.py
+++ b/test/plugins/test_completion.py
@@ -583,9 +583,18 @@ def test_file_completions(workspace, tmpdir):
     # Check completions
     assert len(completions) == 2
     assert [c["kind"] == lsp.CompletionItemKind.File for c in completions]
-    assert (
-        completions[0]["insertText"] == ("bar" + "\\\\")
-        if os.name == "nt"
-        else ("bar" + "\\/")
+    assert completions[0]["insertText"] == (
+        ("bar" + "\\") if os.name == "nt" else ("bar" + "/")
+    )
+    assert completions[1]["insertText"] == 'foo.txt"'
+
+    # When snippets are supported, ensure that path separators are escaped.
+    support_snippet = {
+        "textDocument": {"completion": {"completionItem": {"snippetSupport": True}}}
+    }
+    doc._config.capabilities.update(support_snippet)
+    completions = pylsp_jedi_completions(doc._config, doc, com_position)
+    assert completions[0]["insertText"] == (
+        ("bar" + "\\\\") if os.name == "nt" else ("bar" + "\\/")
     )
     assert completions[1]["insertText"] == 'foo.txt"'
diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py
index c63d279..f97088b 100644
--- a/test/test_notebook_document.py
+++ b/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",
+            },
+        ],
+    }