Add Jupyter notebook cell support via LSP 3.17 Notebook Document Sync (#742)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: edvilme

bundled/tool/lsp_server.py

@@ -15,6 +15,7 @@
 import sysconfig
 import traceback
 from typing import Any, Callable, Dict, List, Optional, Sequence, Union
+from urllib.parse import urlparse, urlunparse
 
 
 # **********************************************************
@@ -74,11 +75,55 @@ def update_environ_path() -> None:
 RUNNER = pathlib.Path(__file__).parent / "runner.py"
 
 MAX_WORKERS = 5
+NOTEBOOK_SYNC_OPTIONS = lsp.NotebookDocumentSyncOptions(
+    notebook_selector=[
+        lsp.NotebookDocumentFilterWithNotebook(
+            notebook="jupyter-notebook",
+            cells=[
+                lsp.NotebookCellLanguage(language="python"),
+            ],
+        ),
+        lsp.NotebookDocumentFilterWithNotebook(
+            notebook="interactive",
+            cells=[
+                lsp.NotebookCellLanguage(language="python"),
+            ],
+        ),
+    ],
+    save=True,
+)
 LSP_SERVER = LanguageServer(
-    name="pylint-server", version="v0.1.0", max_workers=MAX_WORKERS
+    name="pylint-server",
+    version="v0.1.0",
+    max_workers=MAX_WORKERS,
+    notebook_document_sync=NOTEBOOK_SYNC_OPTIONS,
 )
 
 
+def _get_document_path(document: str) -> str:
+    """Returns the filesystem path for a document.
+
+    Examples:
+        file:///path/to/file.py -> /path/to/file.py
+        vscode-notebook-cell:/path/to/notebook.ipynb#C00001 -> /path/to/notebook.ipynb
+    """
+    if not document.startswith("file:"):
+        parsed = urlparse(document)
+        file_uri = urlunparse(
+            (
+                "file",
+                parsed.netloc,
+                parsed.path,
+                parsed.params,
+                parsed.query,
+                "",
+            )
+        )
+        if result := uris.to_fs_path(file_uri):
+            return result
+    return uris.to_fs_path(document) or document
+
+
 # **********************************************************
 # Tool specific code goes below this.
 # **********************************************************
@@ -143,8 +188,92 @@ def did_change(params: lsp.DidChangeTextDocumentParams) -> None:
         )
 
 
-def _linting_helper(document: workspace.TextDocument) -> list[lsp.Diagnostic]:
+@LSP_SERVER.feature(lsp.NOTEBOOK_DOCUMENT_DID_OPEN)
+def notebook_did_open(params: lsp.DidOpenNotebookDocumentParams) -> None:
+    """Run diagnostics on each code cell when a notebook is opened."""
+    nb = LSP_SERVER.workspace.get_notebook_document(
+        notebook_uri=params.notebook_document.uri
+    )
+    if nb is None:
+        return
+    for cell in nb.cells:
+        if cell.kind == lsp.NotebookCellKind.Code and cell.document is not None:
+            _lint_notebook_cell(cell.document)
+
+
+@LSP_SERVER.feature(lsp.NOTEBOOK_DOCUMENT_DID_CHANGE)
+def notebook_did_change(params: lsp.DidChangeNotebookDocumentParams) -> None:
+    """Re-lint cells whose text changed or that were newly added."""
+    if params.change is None or params.change.cells is None:
+        return
+
+    for cell_content in params.change.cells.text_content or []:
+        if cell_content.document:
+            _lint_notebook_cell(cell_content.document.uri)
+
+    structure = params.change.cells.structure
+    if structure and structure.did_open:
+        for cell_document in structure.did_open:
+            _lint_notebook_cell(cell_document.uri)
+
+    if structure and structure.did_close:
+        for cell_document in structure.did_close:
+            _clear_notebook_cell_diagnostics(cell_document.uri)
+
+
+@LSP_SERVER.feature(lsp.NOTEBOOK_DOCUMENT_DID_SAVE)
+def notebook_did_save(params: lsp.DidSaveNotebookDocumentParams) -> None:
+    """Re-lint all cells when a notebook is saved."""
+    nb = LSP_SERVER.workspace.get_notebook_document(
+        notebook_uri=params.notebook_document.uri
+    )
+    if nb is None:
+        return
+    for cell in nb.cells:
+        if cell.kind == lsp.NotebookCellKind.Code and cell.document is not None:
+            _lint_notebook_cell(cell.document)
+
+
+@LSP_SERVER.feature(lsp.NOTEBOOK_DOCUMENT_DID_CLOSE)
+def notebook_did_close(params: lsp.DidCloseNotebookDocumentParams) -> None:
+    """Clear diagnostics for all cells when the notebook is closed."""
+    for cell_doc in params.cell_text_documents:
+        _clear_notebook_cell_diagnostics(cell_doc.uri)
+
+
+def _lint_notebook_cell(cell_uri: str) -> None:
+    """Lint a single notebook cell and publish its diagnostics."""
+    document = LSP_SERVER.workspace.get_text_document(cell_uri)
+    if document is None:
+        return
+    # Update path as pygls generates an invalid path.
+    # TODO: Remove when fixed. # pylint: disable=fixme
+    document.path = _get_document_path(cell_uri)
+    # Linting is only supported for python cells in notebooks.
+    if document.language_id != "python":
+        return
+    diagnostics: list[lsp.Diagnostic] = _linting_helper(document, is_notebook=True)
+    LSP_SERVER.text_document_publish_diagnostics(
+        lsp.PublishDiagnosticsParams(uri=cell_uri, diagnostics=diagnostics)
+    )
+
+
+def _clear_notebook_cell_diagnostics(cell_uri: str) -> None:
+    """Clear diagnostics for a single notebook cell."""
+    LSP_SERVER.text_document_publish_diagnostics(
+        lsp.PublishDiagnosticsParams(uri=cell_uri, diagnostics=[])
+    )
+
+
+def _linting_helper(
+    document: workspace.TextDocument, is_notebook: bool = False
+) -> list[lsp.Diagnostic]:
     try:
+        # Skip notebook cells — they are linted via _lint_notebook_cell which
+        # passes cell content directly, not the notebook file path.
+        if not is_notebook and str(document.uri).startswith("vscode-notebook-cell"):
+            return []
+
         # Notify the client that linting has started for this document.
         LSP_SERVER.protocol.notify(
             "pylint/lintingStarted",
@@ -776,10 +905,6 @@ def _run_tool_on_document(
         log_warning("See `pylint.enabled` in settings.json to enabling linting.")
         return None
 
-    if str(document.uri).startswith("vscode-notebook-cell"):
-        log_warning(f"Skipping notebook cells [Not Supported]: {str(document.uri)}")
-        return None
-
     if utils.is_stdlib_file(document.path):
         log_warning(
             f"Skipping standard library file (stdlib excluded): {document.path}"

package.json

@@ -47,7 +47,9 @@
     "activationEvents": [
         "onLanguage:python",
         "workspaceContains:.pylintrc",
-        "workspaceContains:*.py"
+        "workspaceContains:*.py",
+        "onNotebook:jupyter-notebook",
+        "onNotebook:interactive"
     ],
     "main": "./dist/extension.js",
     "l10n": "./l10n",

src/test/python_tests/lsp_test_client/session.py

@@ -144,6 +144,22 @@ def notify_did_close(self, did_close_params):
         """Sends did close notification to LSP Server."""
         self._send_notification("textDocument/didClose", params=did_close_params)
 
+    def notify_notebook_did_open(self, params):
+        """Sends notebookDocument/didOpen notification to LSP Server."""
+        self._send_notification("notebookDocument/didOpen", params=params)
+
+    def notify_notebook_did_change(self, params):
+        """Sends notebookDocument/didChange notification to LSP Server."""
+        self._send_notification("notebookDocument/didChange", params=params)
+
+    def notify_notebook_did_save(self, params):
+        """Sends notebookDocument/didSave notification to LSP Server."""
+        self._send_notification("notebookDocument/didSave", params=params)
+
+    def notify_notebook_did_close(self, params):
+        """Sends notebookDocument/didClose notification to LSP Server."""
+        self._send_notification("notebookDocument/didClose", params=params)
+
     def text_document_formatting(self, formatting_params):
         """Sends text document format request to LSP server."""
         fut = self._send_request("textDocument/formatting", params=formatting_params)

src/test/python_tests/test_data/sample1/sample_notebook.ipynb

@@ -0,0 +1,39 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "01a29052",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = 1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "36edf88f",
+   "metadata": {},
+   "source": [
+    "# heading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fe27813e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "y = 2"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/test/python_tests/test_data/sample1/sample_notebook_err.ipynb

@@ -0,0 +1,39 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b95324b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3f4c8e46",
+   "metadata": {},
+   "source": [
+    "# heading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "30277226",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "y = 2"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/test/python_tests/test_get_cwd.py

@@ -23,10 +23,10 @@ def feature(self, *args, **kwargs):
         def command(self, *args, **kwargs):
             return lambda f: f
 
-        def show_message_log(self, *args, **kwargs):
+        def window_log_message(self, *args, **kwargs):
             pass
 
-        def show_message(self, *args, **kwargs):
+        def window_show_message(self, *args, **kwargs):
             pass
 
     mock_server = types.ModuleType("pygls.server")
@@ -56,6 +56,10 @@ def show_message(self, *args, **kwargs):
         "INITIALIZE",
         "EXIT",
         "SHUTDOWN",
+        "NOTEBOOK_DOCUMENT_DID_OPEN",
+        "NOTEBOOK_DOCUMENT_DID_CHANGE",
+        "NOTEBOOK_DOCUMENT_DID_SAVE",
+        "NOTEBOOK_DOCUMENT_DID_CLOSE",
     ]:
         setattr(mock_lsp, _name, _name)
     for _name in [
@@ -65,16 +69,28 @@ def show_message(self, *args, **kwargs):
         "DidOpenTextDocumentParams",
         "DidSaveTextDocumentParams",
         "DidChangeTextDocumentParams",
+        "DidChangeNotebookDocumentParams",
+        "DidCloseNotebookDocumentParams",
+        "DidOpenNotebookDocumentParams",
+        "DidSaveNotebookDocumentParams",
         "DocumentFormattingParams",
         "InitializeParams",
+        "NotebookCellKind",
+        "NotebookCellLanguage",
+        "NotebookDocumentFilterWithNotebook",
+        "NotebookDocumentSyncOptions",
         "Position",
         "Range",
         "TextEdit",
         "CodeAction",
         "CodeActionParams",
         "PublishDiagnosticsParams",
     ]:
-        setattr(mock_lsp, _name, type(_name, (), {}))
+        setattr(
+            mock_lsp,
+            _name,
+            type(_name, (), {"__init__": lambda self, *args, **kw: None}),
+        )
     mock_lsp.MessageType = type(
         "MessageType", (), {"Log": 4, "Error": 1, "Warning": 2, "Info": 3}
     )