support notebookDocument__did_close

tkrabel-db · tkrabel-db · commit 2dd71654aeda · 2023-06-14T21:00:20.000+02:00
diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
@@ -466,7 +466,6 @@ def folding(self, doc_uri):
     def m_completion_item__resolve(self, **completionItem):
         return self.completion_item_resolve(completionItem)
 
-    # TODO: add m_notebook_document__did_close
     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'],
@@ -476,6 +475,13 @@ def m_notebook_document__did_open(self, notebookDocument=None, cellTextDocuments
             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.
diff --git a/pylsp/workspace.py b/pylsp/workspace.py
@@ -33,9 +33,9 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
-# pylint: disable=too-many-public-methods
 class Workspace:
 
+    # pylint: disable=too-many-public-methods
 
     M_PUBLISH_DIAGNOSTICS = 'textDocument/publishDiagnostics'
     M_PROGRESS = '$/progress'
diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py
@@ -404,3 +404,74 @@ def test_notebook_document__did_change(client_server_pair): # pylint: disable=re
             ),
         ]
         mock_notify.assert_has_calls(expected_call_args)
+
+
+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)
+
+    # 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
