Add Jupyter notebook cell support via LSP 3.17 Notebook Document Sync

Co-authored-by: edvilme <5952839+edvilme@users.noreply.github.com>

Author: Copilot

README.md

@@ -71,6 +71,47 @@ References, to other extension created by our team using the template:
 - Implementation showing how to handle Formatting. [Black Formatter](https://github.com/microsoft/vscode-black-formatter/tree/main/bundled/tool)
 - Implementation showing how to handle Code Actions. [isort](https://github.com/microsoft/vscode-isort/blob/main/bundled/tool)
 
+## Jupyter Notebook Support
+
+This template includes built-in support for linting and formatting Python cells inside Jupyter notebooks (`.ipynb` files) and the VS Code Interactive Window, following the [LSP 3.17 Notebook Document Sync specification](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#notebookDocument_synchronization).
+
+### How it works
+
+The server declares `NotebookDocumentSyncOptions` (in `lsp_server.py`) that tell the client which notebook types and cell languages to synchronize:
+
+```python
+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,
+)
+```
+
+Four notebook lifecycle handlers are registered:
+
+- `notebookDocument/didOpen` — diagnostics are published for every Python code cell when a notebook is opened.
+- `notebookDocument/didChange` — diagnostics are updated for cells whose text changed, new cells are linted, and removed cells have their diagnostics cleared.
+- `notebookDocument/didSave` — all Python code cells are re-linted on save.
+- `notebookDocument/didClose` — diagnostics are cleared for all cells when the notebook is closed.
+
+The `_get_document_path` helper resolves `vscode-notebook-cell:` URIs back to the parent notebook's filesystem path, so your tool receives a valid path even when processing a cell document.
+
+### Customizing notebook support
+
+- To disable notebook support entirely, remove the `notebook_document_sync=NOTEBOOK_SYNC_OPTIONS` argument from the `LanguageServer` constructor and delete the four `notebook_did_*` handlers.
+- To restrict which cell languages are supported, update the `cells` list in `NOTEBOOK_SYNC_OPTIONS`.
+- To change the linting/formatting behavior per cell, update the individual handlers.
+
+See [microsoft/vscode-isort#565](https://github.com/microsoft/vscode-isort/pull/565) for a reference implementation of notebook support in a production extension.
+
 ## Building and Run the extension
 
 Run the `Debug Extension and Python` configuration form VS Code. That should build and debug the extension in host window.

bundled/tool/lsp_server.py

@@ -46,9 +46,25 @@ def update_sys_path(path_to_add: str, strategy: str) -> None:
 RUNNER = pathlib.Path(__file__).parent / "lsp_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,
+)
 # TODO: Update the language server name and version.
 LSP_SERVER = server.LanguageServer(
-    name="<pytool-display-name>", version="<server version>", max_workers=MAX_WORKERS
+    name="<pytool-display-name>",
+    version="<server version>",
+    max_workers=MAX_WORKERS,
+    notebook_document_sync=NOTEBOOK_SYNC_OPTIONS,
 )
 
 
@@ -112,6 +128,103 @@ def did_close(params: lsp.DidCloseTextDocumentParams) -> None:
     LSP_SERVER.publish_diagnostics(document.uri, [])
 
 
+@LSP_SERVER.feature(lsp.NOTEBOOK_DOCUMENT_DID_OPEN)
+def notebook_did_open(params: lsp.DidOpenNotebookDocumentParams) -> None:
+    """LSP handler for notebookDocument/didOpen request."""
+    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 or cell.document is None:
+            continue
+        document = LSP_SERVER.workspace.get_text_document(cell.document)
+        diagnostics: list[lsp.Diagnostic] = _linting_helper(document)
+        LSP_SERVER.text_document_publish_diagnostics(
+            lsp.PublishDiagnosticsParams(uri=document.uri, diagnostics=diagnostics)
+        )
+
+
+@LSP_SERVER.feature(lsp.NOTEBOOK_DOCUMENT_DID_CHANGE)
+def notebook_did_change(params: lsp.DidChangeNotebookDocumentParams) -> None:
+    """LSP handler for notebookDocument/didChange request."""
+    nb = LSP_SERVER.workspace.get_notebook_document(
+        notebook_uri=params.notebook_document.uri
+    )
+    if nb is None:
+        return
+
+    change = params.change
+    # Re-lint cells whose text content changed.
+    if change.cells and change.cells.text_content:
+        for text_change in change.cells.text_content:
+            document = LSP_SERVER.workspace.get_text_document(
+                text_change.document.uri
+            )
+            diagnostics: list[lsp.Diagnostic] = _linting_helper(document)
+            LSP_SERVER.text_document_publish_diagnostics(
+                lsp.PublishDiagnosticsParams(uri=document.uri, diagnostics=diagnostics)
+            )
+
+    # Lint newly added cells.
+    if change.cells and change.cells.structure and change.cells.structure.did_open:
+        for cell_doc in change.cells.structure.did_open:
+            document = LSP_SERVER.workspace.get_text_document(cell_doc.uri)
+            diagnostics = _linting_helper(document)
+            LSP_SERVER.text_document_publish_diagnostics(
+                lsp.PublishDiagnosticsParams(uri=document.uri, diagnostics=diagnostics)
+            )
+
+    # Clear diagnostics for removed cells.
+    if change.cells and change.cells.structure and change.cells.structure.did_close:
+        for cell_doc in change.cells.structure.did_close:
+            LSP_SERVER.text_document_publish_diagnostics(
+                lsp.PublishDiagnosticsParams(uri=cell_doc.uri, diagnostics=[])
+            )
+
+
+@LSP_SERVER.feature(lsp.NOTEBOOK_DOCUMENT_DID_SAVE)
+def notebook_did_save(params: lsp.DidSaveNotebookDocumentParams) -> None:
+    """LSP handler for notebookDocument/didSave request."""
+    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 or cell.document is None:
+            continue
+        document = LSP_SERVER.workspace.get_text_document(cell.document)
+        diagnostics: list[lsp.Diagnostic] = _linting_helper(document)
+        LSP_SERVER.text_document_publish_diagnostics(
+            lsp.PublishDiagnosticsParams(uri=document.uri, diagnostics=diagnostics)
+        )
+
+
+@LSP_SERVER.feature(lsp.NOTEBOOK_DOCUMENT_DID_CLOSE)
+def notebook_did_close(params: lsp.DidCloseNotebookDocumentParams) -> None:
+    """LSP handler for notebookDocument/didClose request."""
+    for cell_doc in params.cell_text_documents:
+        LSP_SERVER.text_document_publish_diagnostics(
+            lsp.PublishDiagnosticsParams(uri=cell_doc.uri, diagnostics=[])
+        )
+
+
+def _get_document_path(document: workspace.Document) -> str:
+    """Returns the file path for a document, handling notebook cell URIs.
+
+    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.uri.startswith("file:"):
+        return uris.to_fs_path(
+            document.uri.split("#")[0].replace("vscode-notebook-cell:", "file:", 1)
+        )
+    return uris.to_fs_path(document.uri)
+
+
 def _linting_helper(document: workspace.Document) -> list[lsp.Diagnostic]:
     # TODO: Determine if your tool supports passing file content via stdin.
     # If you want to support linting on change then your tool will need to
@@ -436,12 +549,11 @@ def _run_tool_on_document(
     """
     if extra_args is None:
         extra_args = []
-    if str(document.uri).startswith("vscode-notebook-cell"):
-        # TODO: Decide on if you want to skip notebook cells.
-        # Skip notebook cells
-        return None
+    # TODO: Notebook cells are now supported via the notebookDocument/ handlers.
+    # If you want to customize notebook cell handling, update the notebook handlers above.
 
-    if utils.is_stdlib_file(document.path):
+    document_path = _get_document_path(document)
+    if utils.is_stdlib_file(document_path):
         # TODO: Decide on if you want to skip standard library files.
         # Skip standard library python files.
         return None
@@ -486,7 +598,7 @@ def _run_tool_on_document(
         # set use_stdin to False, or provide path, what ever is appropriate for your tool.
         argv += []
     else:
-        argv += [document.path]
+        argv += [document_path]
 
     if use_path:
         # This mode is used when running executables.

package.json

@@ -45,7 +45,9 @@
     },
     "activationEvents": [
         "onLanguage:python",
-        "workspaceContains:*.py"
+        "workspaceContains:*.py",
+        "onNotebook:jupyter-notebook",
+        "onNotebook:interactive"
     ],
     "main": "./dist/extension.js",
     "scripts": {

src/test/python_tests/lsp_test_client/session.py

@@ -141,6 +141,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 references request to LSP server."""
         fut = self._send_request("textDocument/formatting", params=formatting_params)

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

@@ -0,0 +1,27 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cell1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = 1\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.9.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/test/python_tests/test_get_cwd.py

@@ -35,20 +35,32 @@ def show_message(self, *args, **kwargs):
 
     mock_uris = types.ModuleType("pygls.uris")
     mock_uris.from_fs_path = lambda p: "file://" + p
+    mock_uris.to_fs_path = lambda u: u.replace("file://", "")
 
     mock_lsp = types.ModuleType("lsprotocol.types")
     for _name in [
         "TEXT_DOCUMENT_DID_OPEN", "TEXT_DOCUMENT_DID_SAVE", "TEXT_DOCUMENT_DID_CLOSE",
         "TEXT_DOCUMENT_FORMATTING", "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 [
         "Diagnostic", "DiagnosticSeverity", "DidCloseTextDocumentParams",
         "DidOpenTextDocumentParams", "DidSaveTextDocumentParams",
         "DocumentFormattingParams", "InitializeParams", "Position", "Range", "TextEdit",
+        "DidOpenNotebookDocumentParams", "DidChangeNotebookDocumentParams",
+        "DidSaveNotebookDocumentParams", "DidCloseNotebookDocumentParams",
+        "NotebookCellLanguage", "NotebookDocumentFilterWithNotebook",
+        "NotebookDocumentSyncOptions", "PublishDiagnosticsParams",
     ]:
-        setattr(mock_lsp, _name, type(_name, (), {}))
+        def _make_stub(name):
+            def __init__(self, *args, **kwargs):
+                pass
+            return type(name, (), {"__init__": __init__})
+        setattr(mock_lsp, _name, _make_stub(_name))
     mock_lsp.MessageType = type("MessageType", (), {"Log": 4, "Error": 1, "Warning": 2, "Info": 3})
+    mock_lsp.NotebookCellKind = type("NotebookCellKind", (), {"Code": 2})
 
     for _mod_name, _mod in [
         ("pygls", types.ModuleType("pygls")),