Fix notebook path completions root folder (#6018)

Addresses #5948.

Also reworked the tests slightly to allow for testing this a bit more
end-to-end. Now that we're configuring the root path initialization
option in tests, we no longer need to manually change directories.

### Release Notes

#### New Features

- N/A

#### Bug Fixes

- Python notebooks now complete paths in the notebook's parent folder
instead of the workspace root (#5948).

### QA Notes

See the linked issue.

extensions/positron-python/python_files/positron/positron_ipykernel/positron_jedilsp.py

@@ -10,11 +10,12 @@
 import threading
 import warnings
 from functools import lru_cache
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union, cast
 
 from comm.base_comm import BaseComm
 
-from ._vendor import attrs
+from ._vendor import attrs, cattrs
 from ._vendor.jedi.api import Interpreter, Project
 from ._vendor.jedi_language_server import jedi_utils, pygls_utils
 from ._vendor.jedi_language_server.server import (
@@ -35,6 +36,7 @@
 )
 from ._vendor.lsprotocol.types import (
     COMPLETION_ITEM_RESOLVE,
+    INITIALIZE,
     TEXT_DOCUMENT_CODE_ACTION,
     TEXT_DOCUMENT_COMPLETION,
     TEXT_DOCUMENT_DEFINITION,
@@ -69,10 +71,13 @@
     DocumentSymbol,
     DocumentSymbolParams,
     Hover,
+    InitializeParams,
+    InitializeResult,
     InsertTextFormat,
     Location,
     MarkupContent,
     MarkupKind,
+    MessageType,
     Position,
     Range,
     RenameParams,
@@ -87,6 +92,7 @@
 )
 from ._vendor.pygls.capabilities import get_capability
 from ._vendor.pygls.feature_manager import has_ls_param_or_annotation
+from ._vendor.pygls.protocol import lsp_method
 from ._vendor.pygls.workspace.text_document import TextDocument
 from .help_comm import ShowHelpTopicParams
 from .inspectors import BaseColumnInspector, BaseTableInspector, get_inspector
@@ -128,6 +134,15 @@ class HelpTopicRequest:
     jsonrpc: str = attrs.field(default="2.0")
 
 
+@attrs.define
+class PositronInitializationOptions:
+    """
+    Positron-specific language server initialization options.
+    """
+
+    notebook_path: Optional[Path] = attrs.field(default=None)
+
+
 class PositronJediLanguageServerProtocol(JediLanguageServerProtocol):
     @lru_cache()
     def get_message_type(self, method: str) -> Optional[Type]:
@@ -137,6 +152,50 @@ def get_message_type(self, method: str) -> Optional[Type]:
             return HelpTopicRequest
         return super().get_message_type(method)
 
+    @lsp_method(INITIALIZE)
+    def lsp_initialize(self, params: InitializeParams) -> InitializeResult:
+        result = super().lsp_initialize(params)
+
+        server = self._server
+
+        # Parse Positron-specific initialization options.
+        try:
+            raw_initialization_options = (params.initialization_options or {}).get("positron", {})
+            initialization_options = cattrs.structure(
+                raw_initialization_options, PositronInitializationOptions
+            )
+        except cattrs.BaseValidationError as error:
+            # Show an error message in the client.
+            msg = f"Invalid PositronInitializationOptions, using defaults: {cattrs.transform_error(error)}"
+            server.show_message(msg, msg_type=MessageType.Error)
+            server.show_message_log(msg, msg_type=MessageType.Error)
+            initialization_options = PositronInitializationOptions()
+
+        # If a notebook path was provided, set the project path to the notebook's parent.
+        # See /~https://github.com/posit-dev/positron/issues/5948.
+        notebook_path = initialization_options.notebook_path
+        if notebook_path:
+            path = notebook_path.parent
+        else:
+            path = self._server.workspace.root_path
+
+        # Create the Jedi Project.
+        # Note that this overwrites a Project already created in the parent class.
+        workspace_options = server.initialization_options.workspace
+        server.project = (
+            Project(
+                path=path,
+                environment_path=workspace_options.environment_path,
+                added_sys_path=workspace_options.extra_paths,
+                smart_sys_path=True,
+                load_unsafe_extensions=False,
+            )
+            if path
+            else None
+        )
+
+        return result
+
 
 class PositronJediLanguageServer(JediLanguageServer):
     """Positron extension to the Jedi language server."""
@@ -718,7 +777,7 @@ def _publish_diagnostics(server: PositronJediLanguageServer, uri: str) -> None:
     # canceling notifications that happen in that interval.
     # Since this function is executed after a delay, we need to check
     # whether the document still exists
-    if uri not in server.workspace.documents:
+    if uri not in server.workspace.text_documents:
         return
 
     doc = server.workspace.get_text_document(uri)

extensions/positron-python/python_files/positron/positron_ipykernel/tests/test_positron_jedilsp.py

@@ -4,6 +4,8 @@
 #
 
 import os
+from functools import partial
+from pathlib import Path
 from threading import Timer
 from typing import Any, Dict, List, Optional, cast
 from unittest.mock import Mock, patch
@@ -12,6 +14,7 @@
 import polars as pl
 import pytest
 
+from positron_ipykernel._vendor import cattrs
 from positron_ipykernel._vendor.jedi_language_server import jedi_utils
 from positron_ipykernel._vendor.lsprotocol.types import (
     ClientCapabilities,
@@ -30,10 +33,12 @@
     TextDocumentItem,
     TextEdit,
 )
+from positron_ipykernel._vendor.pygls.workspace.text_document import TextDocument
 from positron_ipykernel.help_comm import ShowHelpTopicParams
 from positron_ipykernel.jedi import PositronInterpreter
 from positron_ipykernel.positron_jedilsp import (
     HelpTopicParams,
+    PositronInitializationOptions,
     PositronJediLanguageServer,
     PositronJediLanguageServerProtocol,
     _clear_diagnostics_debounced,
@@ -55,7 +60,11 @@ def _reduce_debounce_time(monkeypatch):
     monkeypatch.setattr(_publish_diagnostics_debounced, "interval_s", 0.05)
 
 
-def create_server(namespace: Optional[Dict[str, Any]] = None) -> PositronJediLanguageServer:
+def create_server(
+    namespace: Optional[Dict[str, Any]] = None,
+    root_path: Optional[Path] = None,
+    notebook_path: Optional[Path] = None,
+) -> PositronJediLanguageServer:
     # Create a server.
     server = PositronJediLanguageServer(
         name="test-server",
@@ -70,13 +79,22 @@ def create_server(namespace: Optional[Dict[str, Any]] = None) -> PositronJediLan
                 text_document=TextDocumentClientCapabilities(
                     completion=CompletionClientCapabilities(
                         completion_item=CompletionClientCapabilitiesCompletionItemType(
+                            # We test markdown docs exclusively.
                             documentation_format=[MarkupKind.Markdown]
                         ),
                     )
                 )
             ),
+            # Optionally set the root path. This seems to only change file completions.
+            root_path=str(root_path) if root_path else None,
             initialization_options={
-                "markup_kind_preferred": MarkupKind.Markdown,
+                # Pass Positron-specific initialization options in serialized format
+                # to test deserialization too.
+                "positron": cattrs.unstructure(
+                    PositronInitializationOptions(
+                        notebook_path=notebook_path,
+                    )
+                ),
             },
         )
     )
@@ -129,17 +147,13 @@ def prop(self) -> str:
 
 
 def _completions(
-    source: str,
-    namespace: Dict[str, Any],
+    server: PositronJediLanguageServer,
+    text_document: TextDocument,
     character: Optional[int] = None,
 ) -> List[CompletionItem]:
-    lines = source.splitlines()
-    line = len(lines) - 1
+    line = len(text_document.lines) - 1
     if character is None:
-        character = len(lines[line])
-    server = create_server(namespace)
-    text_document = create_text_document(server, "file:///foo.py", source)
-
+        character = len(text_document.lines[line])
     params = CompletionParams(TextDocumentIdentifier(text_document.uri), Position(line, character))
     completion_list = positron_completion(server, params)
 
@@ -172,7 +186,9 @@ def test_positron_completion_exact(
     namespace: Dict[str, Any],
     expected_labels: List[str],
 ) -> None:
-    completions = _completions(source, namespace)
+    server = create_server(namespace)
+    text_document = create_text_document(server, "file:///foo.py", source)
+    completions = _completions(server, text_document)
     completion_labels = [completion.label for completion in completions]
     assert completion_labels == expected_labels
 
@@ -190,11 +206,37 @@ def test_positron_completion_contains(
     namespace: Dict[str, Any],
     expected_label: str,
 ) -> None:
-    completions = _completions(source, namespace)
+    server = create_server(namespace)
+    text_document = create_text_document(server, "file:///foo.py", source)
+    completions = _completions(server, text_document)
     completion_labels = [completion.label for completion in completions]
     assert expected_label in completion_labels
 
 
+def assert_has_path_completion(
+    source: str,
+    expected_completion: str,
+    chars_from_end=1,
+    root_path: Optional[Path] = None,
+    notebook_path: Optional[Path] = None,
+):
+    # Replace separators for testing cross-platform.
+    source = source.replace("/", os.path.sep)
+    expected_completion = expected_completion.replace("/", os.path.sep)
+
+    server = create_server(root_path=root_path, notebook_path=notebook_path)
+    text_document = create_text_document(server, "file:///foo.py", source)
+    character = len(source) - chars_from_end
+    completions = _completions(server, text_document, character)
+
+    assert len(completions) == 1
+
+    expected_position = Position(0, character)
+    assert completions[0].text_edit == TextEdit(
+        Range(expected_position, expected_position), expected_completion
+    )
+
+
 def test_path_completion(tmp_path) -> None:
     # See /~https://github.com/posit-dev/positron/issues/5193.
 
@@ -204,43 +246,25 @@ def test_path_completion(tmp_path) -> None:
     file = dir_ / "weather-report.ipynb"
     file.write_text("")
 
-    cwd = os.getcwd()
-
-    def assert_has_path_completion(source: str, completion: str, chars_from_end=1):
-        # Replace separators for testing cross-platform.
-        source = source.replace("/", os.path.sep)
-        completion = completion.replace("/", os.path.sep)
+    _assert_has_path_completion = partial(assert_has_path_completion, root_path=tmp_path)
 
-        character = len(source) - chars_from_end
-        completions = _completions(source, {}, character)
-        assert len(completions) == 1
-        expected_position = Position(0, character)
-        assert completions[0].text_edit == TextEdit(
-            Range(expected_position, expected_position), completion
-        )
+    # Check directory completions at various points around symbols.
+    _assert_has_path_completion('""', f"my-notebooks.new/")
+    # Quotes aren't automatically closed for directories, since the user may want a file.
+    _assert_has_path_completion('"', "my-notebooks.new/", 0)
+    _assert_has_path_completion('"my"', "-notebooks.new/")
+    _assert_has_path_completion('"my-notebooks"', ".new/")
+    _assert_has_path_completion('"my-notebooks."', "new/")
+    _assert_has_path_completion('"my-notebooks.new"', "/")
 
-    try:
-        os.chdir(tmp_path)
-
-        # Check directory completions at various points around symbols.
-        assert_has_path_completion('""', f"my-notebooks.new/")
-        # Quotes aren't automatically closed for directories, since the user may want a file.
-        assert_has_path_completion('"', "my-notebooks.new/", 0)
-        assert_has_path_completion('"my"', "-notebooks.new/")
-        assert_has_path_completion('"my-notebooks"', ".new/")
-        assert_has_path_completion('"my-notebooks."', "new/")
-        assert_has_path_completion('"my-notebooks.new"', "/")
-
-        # Check file completions at various points around symbols.
-        assert_has_path_completion('"my-notebooks.new/"', "weather-report.ipynb")
-        # Quotes are automatically closed for files, since they end the completion.
-        assert_has_path_completion('"my-notebooks.new/', 'weather-report.ipynb"', 0)
-        assert_has_path_completion('"my-notebooks.new/weather"', "-report.ipynb")
-        assert_has_path_completion('"my-notebooks.new/weather-report"', ".ipynb")
-        assert_has_path_completion('"my-notebooks.new/weather-report."', "ipynb")
-        assert_has_path_completion('"my-notebooks.new/weather-report.ipynb"', "")
-    finally:
-        os.chdir(cwd)
+    # Check file completions at various points around symbols.
+    _assert_has_path_completion('"my-notebooks.new/"', "weather-report.ipynb")
+    # Quotes are automatically closed for files, since they end the completion.
+    _assert_has_path_completion('"my-notebooks.new/', 'weather-report.ipynb"', 0)
+    _assert_has_path_completion('"my-notebooks.new/weather"', "-report.ipynb")
+    _assert_has_path_completion('"my-notebooks.new/weather-report"', ".ipynb")
+    _assert_has_path_completion('"my-notebooks.new/weather-report."', "ipynb")
+    _assert_has_path_completion('"my-notebooks.new/weather-report.ipynb"', "")
 
 
 _pd_df = pd.DataFrame({"a": [0]})
@@ -413,3 +437,20 @@ def test_close_notebook_cell_clears_diagnostics():
         for timer in timers:
             timer.join()
         mock.assert_called_once_with(params.text_document.uri, [])
+
+
+def test_notebook_path_completions(tmp_path):
+    # Notebook path completions should be in the notebook's parent, not root path.
+    # See: /~https://github.com/posit-dev/positron/issues/5948
+    notebook_parent = tmp_path / "notebooks"
+    notebook_parent.mkdir()
+
+    notebook_path = notebook_parent / "notebook.ipynb"
+
+    # Create a file in the notebook's parent.
+    file_to_complete = notebook_parent / "data.csv"
+    file_to_complete.write_text("")
+
+    assert_has_path_completion(
+        '""', file_to_complete.name, root_path=tmp_path, notebook_path=notebook_path
+    )

extensions/positron-python/src/client/positron/lsp.ts

@@ -107,6 +107,14 @@ export class PythonLsp implements vscode.Disposable {
         // Override default output channel with our persistant one that is reused across sessions.
         this._clientOptions.outputChannel = outputChannel;
 
+        // Set Positron-specific server initialization options.
+        // If this server is for a notebook, set the notebook path option.
+        if (notebookUri) {
+            this._clientOptions.initializationOptions.positron = {
+                notebook_path: notebookUri.fsPath,
+            };
+        }
+
         const message = `Creating Positron Python ${this._version} language client (port ${port})`;
         traceInfo(message);
         outputChannel.appendLine(message);