Add Jupyter notebook cell support via LSP 3.17 Notebook Document Sync

bundled/tool/lsp_server.py

@@ -15,6 +15,7 @@
 import uuid
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Sequence
+from urllib.parse import urlparse, urlunparse
 
 
 # **********************************************************
@@ -75,11 +76,49 @@ def update_environ_path() -> None:
 GLOBAL_SETTINGS = {}
 
 MAX_WORKERS = 5
-LSP_SERVER = LanguageServer(name="Mypy", version="v0.1.0", max_workers=MAX_WORKERS)
+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="Mypy",
+    version="v0.1.0",
+    max_workers=MAX_WORKERS,
+    notebook_document_sync=NOTEBOOK_SYNC_OPTIONS,
+)
 
 DMYPY_ARGS = {}
 DMYPY_STATUS_FILE_ROOT = None
 
+
+def _get_document_path(document: TextDocument) -> 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.uri.startswith("file:"):
+        parsed = urlparse(document.uri)
+        file_uri = urlunparse(("file", *parsed[1:-1], ""))
+        if result := uris.to_fs_path(file_uri):
+            return result
+    return document.path
+
+
 # **********************************************************
 # Tool specific code goes below this.
 # **********************************************************
@@ -183,6 +222,70 @@ def did_close(params: lsp.DidCloseTextDocumentParams) -> None:
         _clear_diagnostics(document)
 
 
+@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 or cell.document is None:
+            continue
+        document = LSP_SERVER.workspace.get_text_document(cell.document)
+        _linting_helper(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 []:
+        document = LSP_SERVER.workspace.get_text_document(cell_content.document.uri)
+        _linting_helper(document)
+
+    structure = params.change.cells.structure
+    if structure and structure.did_open:
+        for cell_doc in structure.did_open:
+            if cell_doc.language_id != "python":
+                continue
+            document = LSP_SERVER.workspace.get_text_document(cell_doc.uri)
+            _linting_helper(document)
+
+    if structure and structure.did_close:
+        for cell_doc in 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:
+    """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 or cell.document is None:
+            continue
+        document = LSP_SERVER.workspace.get_text_document(cell.document)
+        _linting_helper(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:
+        LSP_SERVER.text_document_publish_diagnostics(
+            lsp.PublishDiagnosticsParams(uri=cell_doc.uri, diagnostics=[])
+        )
+
+
 def _is_empty_diagnostics(
     filepath: str, results: Optional[Dict[str, str | None]]
 ) -> bool:
@@ -210,12 +313,6 @@ def _linting_helper(document: TextDocument) -> None:
         # deep copy here to prevent accidentally updating global settings.
         settings = copy.deepcopy(_get_settings_by_document(document))
 
-        if str(document.uri).startswith("vscode-notebook-cell"):
-            # We don't support running mypy on notebook cells.
-            log_warning(f"Skipping notebook cells [Not Supported]: {str(document.uri)}")
-            _clear_diagnostics(document)
-            return None
-
         if settings["reportingScope"] == "file" and utils.is_stdlib_file(document.path):
             log_warning(
                 f"Skipping standard library file (stdlib excluded): {document.path}"
@@ -290,6 +387,7 @@ def _linting_helper(document: TextDocument) -> None:
                 message=f"Linting failed with error:\r\n{traceback.format_exc()}",
             )
         )
+        _clear_diagnostics(document)
     return []
 
 

package.json

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

src/test/python_tests/lsp_test_client/session.py

@@ -142,6 +142,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.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

@@ -25,13 +25,10 @@ def feature(self, *args, **kwargs):
         def command(self, *args, **kwargs):
             return lambda f: f
 
-        def show_message_log(self, *args, **kwargs):
-            pass
-
-        def show_message(self, *args, **kwargs):
+        def window_log_message(self, *args, **kwargs):
             pass
 
-        def window_log_message(self, *args, **kwargs):
+        def window_show_message(self, *args, **kwargs):
             pass
 
     mock_server = types.ModuleType("pygls.lsp.server")
@@ -53,6 +50,10 @@ def window_log_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 [
@@ -61,10 +62,19 @@ def window_log_message(self, *args, **kwargs):
         "DidCloseTextDocumentParams",
         "DidOpenTextDocumentParams",
         "DidSaveTextDocumentParams",
+        "DidChangeNotebookDocumentParams",
+        "DidCloseNotebookDocumentParams",
+        "DidOpenNotebookDocumentParams",
+        "DidSaveNotebookDocumentParams",
         "DocumentFormattingParams",
         "InitializeParams",
         "LogMessageParams",
+        "NotebookCellKind",
+        "NotebookCellLanguage",
+        "NotebookDocumentFilterWithNotebook",
+        "NotebookDocumentSyncOptions",
         "Position",
+        "PublishDiagnosticsParams",
         "Range",
         "TextEdit",
     ]: