Repository: python-lsp/python-lsp-server
Pull request #389 (2023-06-07 09:07:31)

Add notebooks suppport to pylsp
# Jargon

- LC = Language client
- LS = Language server

# What changes are proposed?

Add Notebook Support to LSP for `notebookDocument/didOpen`, `*/didChange`, `*/didClose` messages (#308). Other messages are pending.

# How is this tested?

I added unit tests and introduced a test helper classes that let us send notification messages to the language server and inspect its state. This helps us assert on messages being sent between LC and LS. More specifically, I tested that the new LS sends correct diagnostics on `notebookDocument/didOpen` and `notebookDocument/didChange`, and that the server keeps the workspace in sync. These helpers could also be used to add stronger testing to `textDocument` related features.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f40eb2f..56d2400 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,23 @@
 # History of changes
 
+## Version 1.7.4 (2023/06/29)
+
+### Issues Closed
+
+* [Issue 393](https://github.com/python-lsp/python-lsp-server/issues/393) - Environment path doesn't expand user directory
+
+In this release 1 issue was closed.
+
+### Pull Requests Merged
+
+* [PR 394](https://github.com/python-lsp/python-lsp-server/pull/394) - Resolve homedir references in Jedi environment path, by [@odiroot](https://github.com/odiroot)
+* [PR 381](https://github.com/python-lsp/python-lsp-server/pull/381) - Report progress even when initialization fails, by [@syphar](https://github.com/syphar)
+* [PR 380](https://github.com/python-lsp/python-lsp-server/pull/380) - Fix pylint hang on file with many errors, by [@hetmankp](https://github.com/hetmankp)
+
+In this release 3 pull requests were closed.
+
+----
+
 ## Version 1.7.3 (2023/05/15)
 
 ### Issues Closed
diff --git a/README.md b/README.md
index 27cb468..5016fd6 100644
--- a/README.md
+++ b/README.md
@@ -161,6 +161,24 @@ pip install 'python-lsp-server[websockets]'
 
 ## Development
 
+Dev install
+
+```
+# create conda env
+conda create --name python-lsp-server python=3.8 -y
+conda activate python-lsp-server
+
+pip install ".[all]"
+pip install ".[websockets]"
+```
+
+Run server with ws
+
+```
+pylsp --ws -v    # Info level logging
+pylsp --ws -v -v # Debug level logging
+```
+
 To run the test suite:
 
 ```sh
diff --git a/pylsp/lsp.py b/pylsp/lsp.py
index f97a0a2..d5f4737 100644
--- a/pylsp/lsp.py
+++ b/pylsp/lsp.py
@@ -90,3 +90,8 @@ class TextDocumentSyncKind:
     NONE = 0
     FULL = 1
     INCREMENTAL = 2
+
+
+class NotebookCellKind:
+    Markup = 1
+    Code = 2
diff --git a/pylsp/plugins/_rope_task_handle.py b/pylsp/plugins/_rope_task_handle.py
index 7854bb2..ed93a05 100644
--- a/pylsp/plugins/_rope_task_handle.py
+++ b/pylsp/plugins/_rope_task_handle.py
@@ -76,7 +76,9 @@ class PylspTaskHandle(BaseTaskHandle):
         self.observers = []
 
     def create_jobset(self, name="JobSet", count: Optional[int] = None):
-        report_iter = self.workspace.report_progress(name, None, None)
+        report_iter = self.workspace.report_progress(
+            name, None, None, skip_token_initialization=True
+        )
         result = PylspJobSet(count, report_iter)
         self.job_sets.append(result)
         self._inform_observers()
diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
index 43f886c..9302c4d 100644
--- a/pylsp/python_lsp.py
+++ b/pylsp/python_lsp.py
@@ -6,6 +6,8 @@ import logging
 import os
 import socketserver
 import threading
+import uuid
+from typing import List, Dict, Any
 import ujson as json
 
 from pylsp_jsonrpc.dispatchers import MethodDispatcher
@@ -14,7 +16,7 @@ from pylsp_jsonrpc.streams import JsonRpcStreamReader, JsonRpcStreamWriter
 
 from . import lsp, _utils, uris
 from .config import config
-from .workspace import Workspace
+from .workspace import Workspace, Document, Notebook
 from ._version import __version__
 
 log = logging.getLogger(__name__)
@@ -266,6 +268,11 @@ class PythonLSPServer(MethodDispatcher):
                 },
                 'openClose': True,
             },
+            'notebookDocumentSync': {
+                'notebookSelector': {
+                    'cells': [{'language': 'python'}]
+                }
+            },
             'workspace': {
                 'workspaceFolders': {
                     'supported': True,
@@ -375,11 +382,79 @@ class PythonLSPServer(MethodDispatcher):
     def lint(self, doc_uri, is_saved):
         # Since we're debounced, the document may no longer be open
         workspace = self._match_uri_to_workspace(doc_uri)
-        if doc_uri in workspace.documents:
-            workspace.publish_diagnostics(
-                doc_uri,
-                flatten(self._hook('pylsp_lint', doc_uri, is_saved=is_saved))
-            )
+        document_object = workspace.documents.get(doc_uri, None)
+        if isinstance(document_object, Document):
+            self._lint_text_document(doc_uri, workspace, is_saved=is_saved)
+        elif isinstance(document_object, Notebook):
+            self._lint_notebook_document(document_object, workspace)
+
+    def _lint_text_document(self, doc_uri, workspace, is_saved):
+        workspace.publish_diagnostics(
+            doc_uri,
+            flatten(self._hook('pylsp_lint', doc_uri, is_saved=is_saved))
+        )
+
+    def _lint_notebook_document(self, notebook_document, workspace):  # pylint: disable=too-many-locals
+        """
+        Lint a notebook document.
+
+        This is a bit more complicated than linting a text document, because we need to
+        send the entire notebook document to the pylsp_lint hook, but we need to send
+        the diagnostics back to the client on a per-cell basis.
+        """
+
+        # First, we create a temp TextDocument that represents the whole notebook
+        # contents. We'll use this to send to the pylsp_lint hook.
+        random_uri = str(uuid.uuid4())
+
+        # cell_list helps us map the diagnostics back to the correct cell later.
+        cell_list: List[Dict[str, Any]] = []
+
+        offset = 0
+        total_source = ""
+        for cell in notebook_document.cells:
+            cell_uri = cell['document']
+            cell_document = workspace.get_cell_document(cell_uri)
+
+            num_lines = cell_document.line_count
+
+            data = {
+                'uri': cell_uri,
+                'line_start': offset,
+                'line_end': offset + num_lines - 1,
+                'source': cell_document.source
+            }
+
+            cell_list.append(data)
+            if offset == 0:
+                total_source = cell_document.source
+            else:
+                total_source += ("\n" + cell_document.source)
+
+            offset += num_lines
+
+        workspace.put_document(random_uri, total_source)
+
+        try:
+            document_diagnostics = flatten(self._hook('pylsp_lint', random_uri, is_saved=True))
+
+            # Now we need to map the diagnostics back to the correct cell and publish them.
+            # Note: this is O(n*m) in the number of cells and diagnostics, respectively.
+            for cell in cell_list:
+                cell_diagnostics = []
+                for diagnostic in document_diagnostics:
+                    start_line = diagnostic['range']['start']['line']
+                    end_line = diagnostic['range']['end']['line']
+
+                    if start_line > cell['line_end'] or end_line < cell['line_start']:
+                        continue
+                    diagnostic['range']['start']['line'] = start_line - cell['line_start']
+                    diagnostic['range']['end']['line'] = end_line - cell['line_start']
+                    cell_diagnostics.append(diagnostic)
+
+                workspace.publish_diagnostics(cell['uri'], cell_diagnostics)
+        finally:
+            workspace.rm_document(random_uri)
 
     def references(self, doc_uri, position, exclude_declaration):
         return flatten(self._hook(
@@ -399,6 +474,91 @@ class PythonLSPServer(MethodDispatcher):
     def m_completion_item__resolve(self, **completionItem):
         return self.completion_item_resolve(completionItem)
 
+    def m_notebook_document__did_open(self, notebookDocument=None, cellTextDocuments=None, **_kwargs):
+        workspace = self._match_uri_to_workspace(notebookDocument['uri'])
+        workspace.put_notebook_document(
+            notebookDocument['uri'],
+            notebookDocument['notebookType'],
+            cells=notebookDocument['cells'],
+            version=notebookDocument.get('version'),
+            metadata=notebookDocument.get('metadata')
+        )
+        for cell in (cellTextDocuments or []):
+            workspace.put_cell_document(cell['uri'], cell['languageId'], cell['text'], version=cell.get('version'))
+        self.lint(notebookDocument['uri'], is_saved=True)
+
+    def m_notebook_document__did_close(self, notebookDocument=None, cellTextDocuments=None, **_kwargs):
+        workspace = self._match_uri_to_workspace(notebookDocument['uri'])
+        for cell in (cellTextDocuments or []):
+            workspace.publish_diagnostics(cell['uri'], [])
+            workspace.rm_document(cell['uri'])
+        workspace.rm_document(notebookDocument['uri'])
+
+    def m_notebook_document__did_change(self, notebookDocument=None, change=None, **_kwargs):
+        """
+        Changes to the notebook document.
+
+        This could be one of the following:
+        1. Notebook metadata changed
+        2. Cell(s) added
+        3. Cell(s) deleted
+        4. Cell(s) data changed
+            4.1 Cell metadata changed
+            4.2 Cell source changed
+        """
+        workspace = self._match_uri_to_workspace(notebookDocument['uri'])
+
+        if change.get('metadata'):
+            # Case 1
+            workspace.update_notebook_metadata(notebookDocument['uri'], change.get('metadata'))
+
+        cells = change.get('cells')
+        if cells:
+            # Change to cells
+            structure = cells.get('structure')
+            if structure:
+                # Case 2 or 3
+                notebook_cell_array_change = structure['array']
+                start = notebook_cell_array_change['start']
+                cell_delete_count = notebook_cell_array_change['deleteCount']
+                if cell_delete_count == 0:
+                    # Case 2
+                    # Cell documents
+                    for cell_document in structure['didOpen']:
+                        workspace.put_cell_document(
+                            cell_document['uri'],
+                            cell_document['languageId'],
+                            cell_document['text'],
+                            cell_document.get('version')
+                        )
+                    # Cell metadata which is added to Notebook
+                    workspace.add_notebook_cells(notebookDocument['uri'], notebook_cell_array_change['cells'], start)
+                else:
+                    # Case 3
+                    # Cell documents
+                    for cell_document in structure['didClose']:
+                        workspace.rm_document(cell_document['uri'])
+                        workspace.publish_diagnostics(cell_document['uri'], [])
+                    # Cell metadata which is removed from Notebook
+                    workspace.remove_notebook_cells(notebookDocument['uri'], start, cell_delete_count)
+
+            data = cells.get('data')
+            if data:
+                # Case 4.1
+                for cell in data:
+                    # update NotebookDocument.cells properties
+                    pass
+
+            text_content = cells.get('textContent')
+            if text_content:
+                # Case 4.2
+                for cell in text_content:
+                    cell_uri = cell['document']['uri']
+                    # Even though the protocol says that `changes` is an array, we assume that it's always a single
+                    # element array that contains the last change to the cell source.
+                    workspace.update_document(cell_uri, cell['changes'][0])
+        self.lint(notebookDocument['uri'], is_saved=True)
+
     def m_text_document__did_close(self, textDocument=None, **_kwargs):
         workspace = self._match_uri_to_workspace(textDocument['uri'])
         workspace.publish_diagnostics(textDocument['uri'], [])
diff --git a/pylsp/workspace.py b/pylsp/workspace.py
index ea7f55e..e5fccc9 100644
--- a/pylsp/workspace.py
+++ b/pylsp/workspace.py
@@ -8,7 +8,7 @@ import os
 import re
 import uuid
 import functools
-from typing import Optional, Generator, Callable
+from typing import Optional, Generator, Callable, List
 from threading import RLock
 
 import jedi
@@ -35,6 +35,8 @@ def lock(method):
 
 class Workspace:
 
+    # pylint: disable=too-many-public-methods
+
     M_PUBLISH_DIAGNOSTICS = 'textDocument/publishDiagnostics'
     M_PROGRESS = '$/progress'
     M_INITIALIZE_PROGRESS = 'window/workDoneProgress/create'
@@ -105,12 +107,30 @@ class Workspace:
         """
         return self._docs.get(doc_uri) or self._create_document(doc_uri)
 
+    def get_cell_document(self, doc_uri):
+        return self._docs.get(doc_uri)
+
     def get_maybe_document(self, doc_uri):
         return self._docs.get(doc_uri)
 
     def put_document(self, doc_uri, source, version=None):
         self._docs[doc_uri] = self._create_document(doc_uri, source=source, version=version)
 
+    def put_notebook_document(self, doc_uri, notebook_type, cells, version=None, metadata=None):
+        self._docs[doc_uri] = self._create_notebook_document(doc_uri, notebook_type, cells, version, metadata)
+
+    def add_notebook_cells(self, doc_uri, cells, start):
+        self._docs[doc_uri].add_cells(cells, start)
+
+    def remove_notebook_cells(self, doc_uri, start, delete_count):
+        self._docs[doc_uri].remove_cells(start, delete_count)
+
+    def update_notebook_metadata(self, doc_uri, metadata):
+        self._docs[doc_uri].metadata = metadata
+
+    def put_cell_document(self, doc_uri, language_id, source, version=None):
+        self._docs[doc_uri] = self._create_cell_document(doc_uri, language_id, source, version)
+
     def rm_document(self, doc_uri):
         self._docs.pop(doc_uri)
 
@@ -135,7 +155,20 @@ class Workspace:
         title: str,
         message: Optional[str] = None,
         percentage: Optional[int] = None,
+        skip_token_initialization: bool = False,
     ) -> Generator[Callable[[str, Optional[int]], None], None, None]:
+        """
+        Report progress to the editor / client.
+
+        ``skip_token_initialization` is necessary due to some current
+        limitations of our LSP implementation. When `report_progress`
+        is used from a synchronous LSP handler, the token initialization
+        will time out because we can't receive the response.
+
+        Many editors will still correctly show the progress messages though, which
+        is why we are giving progress users the option to skip the initialization
+        of the progress token.
+        """
         if self._config:
             client_supports_progress_reporting = (
                 self._config.capabilities.get("window", {}).get("workDoneProgress", False)
@@ -144,30 +177,21 @@ class Workspace:
             client_supports_progress_reporting = False
 
         if client_supports_progress_reporting:
-            try:
-                token = self._progress_begin(title, message, percentage)
-            except Exception:  # pylint: disable=broad-exception-caught
-                log.warning(
-                    "There was an error while trying to initialize progress reporting."
-                    "Likely progress reporting was used in a synchronous LSP handler, "
-                    "which is not supported by progress reporting yet.",
-                    exc_info=True
-                )
+            token = self._progress_begin(title, message, percentage, skip_token_initialization)
 
-            else:
-                def progress_message(message: str, percentage: Optional[int] = None) -> None:
-                    self._progress_report(token, message, percentage)
+            def progress_message(message: str, percentage: Optional[int] = None) -> None:
+                self._progress_report(token, message, percentage)
 
-                try:
-                    yield progress_message
-                finally:
-                    self._progress_end(token)
+            try:
+                yield progress_message
+            finally:
+                self._progress_end(token)
 
-                return
+            return
 
         # FALLBACK:
-        # If the client doesn't support progress reporting, or if we failed to
-        # initialize it, we have a dummy method for the caller to use.
+        # If the client doesn't support progress reporting, we have a dummy method
+        # for the caller to use.
         def dummy_progress_message(message: str, percentage: Optional[int] = None) -> None:
             # pylint: disable=unused-argument
             pass
@@ -179,10 +203,23 @@ class Workspace:
         title: str,
         message: Optional[str] = None,
         percentage: Optional[int] = None,
+        skip_token_initialization: bool = False,
     ) -> str:
         token = str(uuid.uuid4())
 
-        self._endpoint.request(self.M_INITIALIZE_PROGRESS, {'token': token}).result(timeout=1.0)
+        if not skip_token_initialization:
+            try:
+                self._endpoint.request(self.M_INITIALIZE_PROGRESS, {'token': token}).result(timeout=1.0)
+            except Exception:  # pylint: disable=broad-exception-caught
+                log.warning(
+                    "There was an error while trying to initialize progress reporting."
+                    "Likely progress reporting was used in a synchronous LSP handler, "
+                    "which is not supported by progress reporting yet. "
+                    "To prevent waiting for the timeout you can set "
+                    "`skip_token_initialization=True`. "
+                    "Not every editor will show progress then, but many will.",
+                    exc_info=True
+                )
 
         value = {
             "kind": "begin",
@@ -258,6 +295,29 @@ class Workspace:
             rope_project_builder=self._rope_project_builder,
         )
 
+    def _create_notebook_document(self, doc_uri, notebook_type, cells, version=None, metadata=None):
+        return Notebook(
+            doc_uri,
+            notebook_type,
+            self,
+            cells=cells,
+            version=version,
+            metadata=metadata
+        )
+
+    def _create_cell_document(self, doc_uri, language_id, source=None, version=None):
+        # TODO: remove what is unnecessary here.
+        path = uris.to_fs_path(doc_uri)
+        return Cell(
+            doc_uri,
+            language_id=language_id,
+            workspace=self,
+            source=source,
+            version=version,
+            extra_sys_path=self.source_roots(path),
+            rope_project_builder=self._rope_project_builder,
+        )
+
     def close(self):
         if self.__rope_autoimport is not None:
             self.__rope_autoimport.close()
@@ -389,6 +449,11 @@ class Document:
             jedi.settings.auto_import_modules = jedi_settings.get('auto_import_modules',
                                                                   DEFAULT_AUTO_IMPORT_MODULES)
             environment_path = jedi_settings.get('environment')
+            # Jedi itself cannot deal with homedir-relative paths.
+            # On systems, where it is expected, expand the home directory.
+            if environment_path and os.name != 'nt':
+                environment_path = os.path.expanduser(environment_path)
+
             extra_paths = jedi_settings.get('extra_paths') or []
             env_vars = jedi_settings.get('env_vars')
 
@@ -441,3 +506,45 @@ class Document:
         environment = self.get_enviroment(environment_path=environment_path, env_vars=env_vars)
         path.extend(environment.get_sys_path())
         return path
+
+
+class Notebook:
+    """Represents a notebook."""
+    def __init__(self, uri, notebook_type, workspace, cells=None, version=None, metadata=None):
+        self.uri = uri
+        self.notebook_type = notebook_type
+        self.workspace = workspace
+        self.version = version
+        self.cells = cells or []
+        self.metadata = metadata or {}
+
+    def __str__(self):
+        return "Notebook with URI '%s'" % str(self.uri)
+
+    def add_cells(self, new_cells: List, start: int) -> None:
+        self.cells[start:start] = new_cells
+
+    def remove_cells(self, start: int, delete_count: int) -> None:
+        del self.cells[start:start+delete_count]
+
+
+class Cell(Document):
+    """
+    Represents a cell in a notebook.
+
+    Notes
+    -----
+    We inherit from Document for now to get the same API. However, a cell document differs from text documents in that
+    they have a language id.
+    """
+
+    def __init__(self, uri, language_id, workspace, source=None, version=None, local=True, extra_sys_path=None,
+                 rope_project_builder=None):
+        super().__init__(uri, workspace, source, version, local, extra_sys_path, rope_project_builder)
+        self.language_id = language_id
+
+    @property
+    @lock
+    def line_count(self):
+        """"Return the number of lines in the cell document."""
+        return len(self.source.split('\n'))


Test patch:
diff --git a/test/data/publish_diagnostics_message_examples/example_1.json b/test/data/publish_diagnostics_message_examples/example_1.json
new file mode 100644
index 0000000..25d43a1
--- /dev/null
+++ b/test/data/publish_diagnostics_message_examples/example_1.json
@@ -0,0 +1,36 @@
+{
+    "diagnostics": [
+        {
+            "message": "invalid syntax",
+            "range": {
+                "end": {
+                    "character": 15,
+                    "line": 1
+                },
+                "start": {
+                    "character": 7,
+                    "line": 1
+                }
+            },
+            "severity": 1,
+            "source": "pyflakes"
+        },
+        {
+            "code": "W292",
+            "message": "W292 no newline at end of file",
+            "range": {
+                "end": {
+                    "character": 7,
+                    "line": 1
+                },
+                "start": {
+                    "character": 7,
+                    "line": 1
+                }
+            },
+            "severity": 2,
+            "source": "pycodestyle"
+        }
+    ],
+    "uri": "/Users/.../code/python-lsp-server/test"
+}
diff --git a/test/data/publish_diagnostics_message_examples/example_2.json b/test/data/publish_diagnostics_message_examples/example_2.json
new file mode 100644
index 0000000..006f95a
--- /dev/null
+++ b/test/data/publish_diagnostics_message_examples/example_2.json
@@ -0,0 +1,68 @@
+{
+    "diagnostics": [
+        {
+            "message": "'sys' imported but unused",
+            "range": {
+                "end": {
+                    "character": 11,
+                    "line": 0
+                },
+                "start": {
+                    "character": 0,
+                    "line": 0
+                }
+            },
+            "severity": 2,
+            "source": "pyflakes"
+        },
+        {
+            "code": "E225",
+            "message": "E225 missing whitespace around operator",
+            "range": {
+                "end": {
+                    "character": 4,
+                    "line": 1
+                },
+                "start": {
+                    "character": 1,
+                    "line": 1
+                }
+            },
+            "severity": 2,
+            "source": "pycodestyle"
+        },
+        {
+            "code": "W292",
+            "message": "W292 no newline at end of file",
+            "range": {
+                "end": {
+                    "character": 5,
+                    "line": 2
+                },
+                "start": {
+                    "character": 5,
+                    "line": 2
+                }
+            },
+            "severity": 2,
+            "source": "pycodestyle"
+        },
+        {
+            "code": "E225",
+            "message": "E225 missing whitespace around operator",
+            "range": {
+                "end": {
+                    "character": 5,
+                    "line": 2
+                },
+                "start": {
+                    "character": 1,
+                    "line": 2
+                }
+            },
+            "severity": 2,
+            "source": "pycodestyle"
+        }
+    ],
+    uri: "/Users/.../code/python-lsp-server/test"
+}
\ No newline at end of file
diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py
new file mode 100644
index 0000000..a85ff93
--- /dev/null
+++ b/test/test_notebook_document.py
@@ -0,0 +1,584 @@
+# Copyright 2021- Python Language Server Contributors.
+
+import os
+import time
+from threading import Thread
+from unittest.mock import patch, call
+
+import pytest
+
+from pylsp import IS_WIN
+from pylsp.python_lsp import PythonLSPServer
+from pylsp.lsp import NotebookCellKind
+
+CALL_TIMEOUT_IN_SECONDS = 30
+
+
+def wait_for_condition(condition, timeout=CALL_TIMEOUT_IN_SECONDS):
+    """Wait for a condition to be true, or timeout."""
+    start_time = time.time()
+    while not condition():
+        time.sleep(0.1)
+        if time.time() - start_time > timeout:
+            raise TimeoutError("Timeout waiting for condition")
+
+
+def start(obj):
+    obj.start()
+
+
+class ClientServerPair:
+    """A class to setup a client/server pair"""
+
+    def __init__(self):
+        # Client to Server pipe
+        csr, csw = os.pipe()
+        # Server to client pipe
+        scr, scw = os.pipe()
+
+        self.server = PythonLSPServer(os.fdopen(csr, "rb"), os.fdopen(scw, "wb"))
+        self.server_thread = Thread(target=start, args=[self.server])
+        self.server_thread.start()
+
+        self.client = PythonLSPServer(os.fdopen(scr, "rb"), os.fdopen(csw, "wb"))
+        self.client_thread = Thread(target=start, args=[self.client])
+        self.client_thread.start()
+
+
+@pytest.fixture
+def client_server_pair():
+    """A fixture that sets up a client/server pair and shuts down the server"""
+    client_server_pair_obj = ClientServerPair()
+
+    yield (client_server_pair_obj.client, client_server_pair_obj.server)
+
+    shutdown_response = client_server_pair_obj.client._endpoint.request(
+        "shutdown"
+    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    assert shutdown_response is None
+    client_server_pair_obj.client._endpoint.notify("exit")
+
+
+@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+def test_initialize(client_server_pair):  # pylint: disable=redefined-outer-name
+    client, server = client_server_pair
+    response = client._endpoint.request(
+        "initialize",
+        {
+            "processId": 1234,
+            "rootPath": os.path.dirname(__file__),
+            "initializationOptions": {},
+        },
+    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    assert server.workspace is not None
+    assert "notebookDocumentSync" in response["capabilities"].keys()
+
+
+@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+def test_notebook_document__did_open(
+    client_server_pair,
+):  # pylint: disable=redefined-outer-name
+    client, server = client_server_pair
+    client._endpoint.request(
+        "initialize",
+        {
+            "processId": 1234,
+            "rootPath": os.path.dirname(__file__),
+            "initializationOptions": {},
+        },
+    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        client._endpoint.notify(
+            "notebookDocument/didOpen",
+            {
+                "notebookDocument": {
+                    "uri": "notebook_uri",
+                    "notebookType": "jupyter-notebook",
+                    "cells": [
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_1_uri",
+                        },
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_2_uri",
+                        },
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_3_uri",
+                        },
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_4_uri",
+                        },
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_5_uri",
+                        },
+                    ],
+                },
+                # Test as many edge cases as possible for the diagnostics message
+                "cellTextDocuments": [
+                    {
+                        "uri": "cell_1_uri",
+                        "languageId": "python",
+                        "text": "",
+                    },
+                    {
+                        "uri": "cell_2_uri",
+                        "languageId": "python",
+                        "text": "\n",
+                    },
+                    {
+                        "uri": "cell_3_uri",
+                        "languageId": "python",
+                        "text": "\nimport sys\n\nabc\n\n",
+                    },
+                    {
+                        "uri": "cell_4_uri",
+                        "languageId": "python",
+                        "text": "x",
+                    },
+                    {
+                        "uri": "cell_5_uri",
+                        "languageId": "python",
+                        "text": "y\n",
+                    },
+                ],
+            },
+        )
+        wait_for_condition(lambda: mock_notify.call_count >= 5)
+        expected_call_args = [
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_1_uri",
+                    "diagnostics": [],
+                },
+            ),
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_2_uri",
+                    "diagnostics": [],
+                },
+            ),
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_3_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pyflakes",
+                            "range": {
+                                "start": {"line": 1, "character": 0},
+                                "end": {"line": 1, "character": 11},
+                            },
+                            "message": "'sys' imported but unused",
+                            "severity": 2,
+                        },
+                        {
+                            "source": "pyflakes",
+                            "range": {
+                                "start": {"line": 3, "character": 0},
+                                "end": {"line": 3, "character": 4},
+                            },
+                            "message": "undefined name 'abc'",
+                            "severity": 1,
+                        },
+                        {
+                            "source": "pycodestyle",
+                            "range": {
+                                "start": {"line": 1, "character": 0},
+                                "end": {"line": 1, "character": 11},
+                            },
+                            "message": "E303 too many blank lines (4)",
+                            "code": "E303",
+                            "severity": 2,
+                        },
+                    ],
+                },
+            ),
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_4_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pyflakes",
+                            "range": {
+                                "start": {"line": 0, "character": 0},
+                                "end": {"line": 0, "character": 2},
+                            },
+                            "message": "undefined name 'x'",
+                            "severity": 1,
+                        },
+                    ],
+                },
+            ),
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_5_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pyflakes",
+                            "range": {
+                                "start": {"line": 0, "character": 0},
+                                "end": {"line": 0, "character": 2},
+                            },
+                            "message": "undefined name 'y'",
+                            "severity": 1,
+                        },
+                    ],
+                },
+            ),
+        ]
+        mock_notify.assert_has_calls(expected_call_args)
+
+
+@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+def test_notebook_document__did_change(
+    client_server_pair,
+):  # pylint: disable=redefined-outer-name
+    client, server = client_server_pair
+    client._endpoint.request(
+        "initialize",
+        {
+            "processId": 1234,
+            "rootPath": os.path.dirname(__file__),
+            "initializationOptions": {},
+        },
+    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+
+    # Open notebook
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        client._endpoint.notify(
+            "notebookDocument/didOpen",
+            {
+                "notebookDocument": {
+                    "uri": "notebook_uri",
+                    "notebookType": "jupyter-notebook",
+                    "cells": [
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_1_uri",
+                        },
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_2_uri",
+                        },
+                    ],
+                },
+                "cellTextDocuments": [
+                    {
+                        "uri": "cell_1_uri",
+                        "languageId": "python",
+                        "text": "import sys",
+                    },
+                    {
+                        "uri": "cell_2_uri",
+                        "languageId": "python",
+                        "text": "",
+                    },
+                ],
+            },
+        )
+        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
+        assert len(server.workspace.get_document("notebook_uri").cells) == 2
+        expected_call_args = [
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_1_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pyflakes",
+                            "range": {
+                                "start": {"line": 0, "character": 0},
+                                "end": {"line": 0, "character": 11},
+                            },
+                            "message": "'sys' imported but unused",
+                            "severity": 2,
+                        }
+                    ],
+                },
+            ),
+            call(
+                "textDocument/publishDiagnostics",
+                params={"uri": "cell_2_uri", "diagnostics": []},
+            ),
+        ]
+        mock_notify.assert_has_calls(expected_call_args)
+
+    # Remove second cell
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        client._endpoint.notify(
+            "notebookDocument/didChange",
+            {
+                "notebookDocument": {
+                    "uri": "notebook_uri",
+                },
+                "change": {
+                    "cells": {
+                        "structure": {
+                            "array": {
+                                "start": 1,
+                                "deleteCount": 1,
+                            },
+                            "didClose": [
+                                {
+                                    "uri": "cell_2_uri",
+                                }
+                            ],
+                        },
+                    }
+                },
+            },
+        )
+        wait_for_condition(lambda: mock_notify.call_count >= 2)
+        assert len(server.workspace.documents) == 2
+        assert "cell_2_uri" not in server.workspace.documents
+        assert len(server.workspace.get_document("notebook_uri").cells) == 1
+        expected_call_args = [
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_1_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pyflakes",
+                            "range": {
+                                "start": {"line": 0, "character": 0},
+                                "end": {"line": 0, "character": 10},
+                            },
+                            "message": "'sys' imported but unused",
+                            "severity": 2,
+                        },
+                        {
+                            "source": "pycodestyle",
+                            "range": {
+                                "start": {"line": 0, "character": 10},
+                                "end": {"line": 0, "character": 10},
+                            },
+                            "message": "W292 no newline at end of file",
+                            "code": "W292",
+                            "severity": 2,
+                        },
+                    ],
+                },
+            )
+        ]
+        mock_notify.assert_has_calls(expected_call_args)
+
+    # Add second cell
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        client._endpoint.notify(
+            "notebookDocument/didChange",
+            {
+                "notebookDocument": {
+                    "uri": "notebook_uri",
+                },
+                "change": {
+                    "cells": {
+                        "structure": {
+                            "array": {
+                                "start": 1,
+                                "deleteCount": 0,
+                                "cells": [
+                                    {
+                                        "kind": NotebookCellKind.Code,
+                                        "document": "cell_3_uri",
+                                    }
+                                ],
+                            },
+                            "didOpen": [
+                                {
+                                    "uri": "cell_3_uri",
+                                    "languageId": "python",
+                                    "text": "x",
+                                }
+                            ],
+                        },
+                    }
+                },
+            },
+        )
+        wait_for_condition(lambda: mock_notify.call_count >= 2)
+        assert len(server.workspace.documents) == 3
+        assert "cell_3_uri" in server.workspace.documents
+        assert len(server.workspace.get_document("notebook_uri").cells) == 2
+        expected_call_args = [
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_1_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pyflakes",
+                            "range": {
+                                "start": {"line": 0, "character": 0},
+                                "end": {"line": 0, "character": 11},
+                            },
+                            "message": "'sys' imported but unused",
+                            "severity": 2,
+                        }
+                    ],
+                },
+            ),
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_3_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pyflakes",
+                            "range": {
+                                "start": {"line": 0, "character": 0},
+                                "end": {"line": 0, "character": 1},
+                            },
+                            "message": "undefined name 'x'",
+                            "severity": 1,
+                        },
+                        {
+                            "source": "pycodestyle",
+                            "range": {
+                                "start": {"line": 0, "character": 1},
+                                "end": {"line": 0, "character": 1},
+                            },
+                            "message": "W292 no newline at end of file",
+                            "code": "W292",
+                            "severity": 2,
+                        },
+                    ],
+                },
+            ),
+        ]
+        mock_notify.assert_has_calls(expected_call_args)
+
+    # Edit second cell
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        client._endpoint.notify(
+            "notebookDocument/didChange",
+            {
+                "notebookDocument": {
+                    "uri": "notebook_uri",
+                },
+                "change": {
+                    "cells": {
+                        "textContent": [
+                            {
+                                "document": {
+                                    "uri": "cell_3_uri",
+                                },
+                                "changes": [{"text": "sys.path"}],
+                            }
+                        ]
+                    }
+                },
+            },
+        )
+        wait_for_condition(lambda: mock_notify.call_count >= 2)
+        expected_call_args = [
+            call(
+                "textDocument/publishDiagnostics",
+                params={"uri": "cell_1_uri", "diagnostics": []},
+            ),
+            call(
+                "textDocument/publishDiagnostics",
+                params={
+                    "uri": "cell_3_uri",
+                    "diagnostics": [
+                        {
+                            "source": "pycodestyle",
+                            "range": {
+                                "start": {"line": 0, "character": 8},
+                                "end": {"line": 0, "character": 8},
+                            },
+                            "message": "W292 no newline at end of file",
+                            "code": "W292",
+                            "severity": 2,
+                        }
+                    ],
+                },
+            ),
+        ]
+        mock_notify.assert_has_calls(expected_call_args)
+
+
+@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+def test_notebook__did_close(
+    client_server_pair,
+):  # pylint: disable=redefined-outer-name
+    client, server = client_server_pair
+    client._endpoint.request(
+        "initialize",
+        {
+            "processId": 1234,
+            "rootPath": os.path.dirname(__file__),
+            "initializationOptions": {},
+        },
+    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+
+    # Open notebook
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        client._endpoint.notify(
+            "notebookDocument/didOpen",
+            {
+                "notebookDocument": {
+                    "uri": "notebook_uri",
+                    "notebookType": "jupyter-notebook",
+                    "cells": [
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_1_uri",
+                        },
+                        {
+                            "kind": NotebookCellKind.Code,
+                            "document": "cell_2_uri",
+                        },
+                    ],
+                },
+                "cellTextDocuments": [
+                    {
+                        "uri": "cell_1_uri",
+                        "languageId": "python",
+                        "text": "import sys",
+                    },
+                    {
+                        "uri": "cell_2_uri",
+                        "languageId": "python",
+                        "text": "",
+                    },
+                ],
+            },
+        )
+        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
+
+    # Close notebook
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        client._endpoint.notify(
+            "notebookDocument/didClose",
+            {
+                "notebookDocument": {
+                    "uri": "notebook_uri",
+                },
+                "cellTextDocuments": [
+                    {
+                        "uri": "cell_1_uri",
+                    },
+                    {
+                        "uri": "cell_2_uri",
+                    },
+                ],
+            },
+        )
+        wait_for_condition(lambda: mock_notify.call_count >= 2)
+        assert len(server.workspace.documents) == 0
diff --git a/test/test_workspace.py b/test/test_workspace.py
index 94a9cba..88b1650 100644
--- a/test/test_workspace.py
+++ b/test/test_workspace.py
@@ -23,6 +23,16 @@ def test_put_document(pylsp):
     assert DOC_URI in pylsp.workspace._docs
 
 
+def test_put_notebook_document(pylsp):
+    pylsp.workspace.put_notebook_document(DOC_URI, 'jupyter-notebook', [])
+    assert DOC_URI in pylsp.workspace._docs
+
+
+def test_put_cell_document(pylsp):
+    pylsp.workspace.put_cell_document(DOC_URI, 'python', 'content')
+    assert DOC_URI in pylsp.workspace._docs
+
+
 def test_get_document(pylsp):
     pylsp.workspace.put_document(DOC_URI, 'TEXT')
     assert pylsp.workspace.get_document(DOC_URI).source == 'TEXT'
@@ -327,19 +337,24 @@ def test_progress_simple(workspace, consumer):
 
 
 @pytest.mark.parametrize("exc", [Exception("something"), TimeoutError()])
-def test_progress_initialization_fails(workspace, consumer, endpoint, exc):
+def test_progress_initialization_fails_but_is_skipped(workspace, consumer, endpoint, exc):
     def failing_token_initialization(self, *_args, **_kwargs):
         raise exc
     endpoint._dispatcher.m_window__work_done_progress__create = failing_token_initialization
 
     workspace._config.capabilities['window'] = {"workDoneProgress": True}
 
-    with workspace.report_progress("some_title"):
+    with workspace.report_progress("some_title", skip_token_initialization=True):
         pass
 
     # we only see the failing token initialization call, no other calls
-    init_call, = consumer.call_args_list
-    assert init_call[0][0]['method'] == 'window/workDoneProgress/create'
+    progress_calls = consumer.call_args_list
+    assert all(call[0][0]["method"] == "$/progress" for call in progress_calls)
+    assert len({call[0][0]["params"]["token"] for call in progress_calls}) == 1
+    assert [call[0][0]["params"]["value"] for call in progress_calls] == [
+        {"kind": "begin", "title": "some_title"},
+        {"kind": "end"},
+    ]
 
 
 def test_progress_with_percent(workspace, consumer):