Add notebooks suppport to pylsp

README.md

@@ -161,6 +161,24 @@ pip install 'python-lsp-server[websockets]'
 
 ## Development
 
+Dev install
+
+```
+# create conda env
+cc python-lsp-server
+ca 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

pylsp/lsp.py

@@ -90,3 +90,8 @@ class TextDocumentSyncKind:
     NONE = 0
     FULL = 1
     INCREMENTAL = 2
+
+
+class NotebookCellKind:
+    Markup = 1
+    Code = 2

pylsp/lsp_types.py

@@ -0,0 +1,61 @@
+from typing import TypedDict, Optional, Union, List, NewType
+from .lsp import DiagnosticSeverity, DiagnosticTag
+
+"""
+Types derived from the LSP protocol
+See: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/
+"""
+
+URI = NewType('URI', str)
+
+NotebookCell = TypedDict('NotebookCell', {
+    'kind': str,
+    'document': URI,
+    'metadata': Optional[dict],
+    'executionSummary': Optional[dict],
+})
+
+NotebookDocument = TypedDict('NotebookDocument', {
+    'uri': str,
+    'notebookType': str,
+    'version': int,
+    'metadata': Optional[dict],
+    'cells': List[NotebookCell],
+})
+
+CodeDescription = TypedDict('CodeDescription', {
+    'href': URI,
+})
+
+Position = TypedDict('Position', {
+    'line': int,
+    'character': int,
+})
+
+Range = TypedDict('Range', {
+    'start': Position,
+    'end': Position,
+})
+
+Location = TypedDict('Location', {
+    'uri': URI,
+    'range': Range,
+})
+
+DiagnosticRelatedInformation = TypedDict('DiagnosticRelatedInformation', {
+    'location': dict,
+    'message': str,
+})
+
+Diagnostic = TypedDict('Diagnostic', {
+    'range': dict,
+    'severity': Optional[DiagnosticSeverity],
+    'code': Optional[Union[int, str]],
+    'codeDescription': Optional[CodeDescription],
+    'source': Optional[str],
+    'message': str,
+    'tags': Optional[List[DiagnosticTag]],
+    'relatedInformation': Optional[List[DiagnosticRelatedInformation]],
+    'data': Optional[dict],
+})
+

pylsp/python_lsp.py

@@ -7,15 +7,20 @@
 import socketserver
 import threading
 import ujson as json
+import uuid
+
 
 from pylsp_jsonrpc.dispatchers import MethodDispatcher
 from pylsp_jsonrpc.endpoint import Endpoint
 from pylsp_jsonrpc.streams import JsonRpcStreamReader, JsonRpcStreamWriter
 
+from typing import List, Dict, Any
+
 from . import lsp, _utils, uris
 from .config import config
-from .workspace import Workspace
+from .workspace import Workspace, Document, Cell, Notebook
 from ._version import __version__
+from .lsp_types import Diagnostic
 
 log = logging.getLogger(__name__)
 
@@ -266,6 +271,10 @@ def capabilities(self):
                 },
                 'openClose': True,
             },
+            # TODO: add notebookDocumentSync when we support it
+            # 'notebookDocumentSync' : {
+            #     'notebook': '*',
+            # },
             'workspace': {
                 'workspaceFolders': {
                     'supported': True,
@@ -375,11 +384,74 @@ def hover(self, doc_uri, position):
     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):
+        """
+        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)
+
+            lines = cell_document.lines
+            num_lines = len(lines)
+            start = offset + 1
+            end = offset + num_lines
+            offset += num_lines
+
+            data = {
+                'uri': cell_uri,
+                'line_start': start,
+                'line_end': end,
+                'source': cell_document.source
+            }
+
+            cell_list.append(data)
+            total_source = total_source + "\n" + cell_document.source
+
+        workspace.put_document(random_uri, total_source) 
+        document_diagnostics: List[Diagnostic] = 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.
+        for cell in cell_list:
+            cell_diagnostics: List[Diagnostic] = []
+            for diagnostic in document_diagnostics:
+                if diagnostic['range']['start']['line'] > cell['line_end']:
+                    break
+                diagnostic['range']['start']['line'] = diagnostic['range']['start']['line'] - cell['line_start'] + 1
+                diagnostic['range']['end']['line'] = diagnostic['range']['end']['line'] - cell['line_start'] + 1
+                cell_diagnostics.append(diagnostic)
+                document_diagnostics.pop(0)
+
+            workspace.publish_diagnostics(cell['uri'], cell_diagnostics)
+
+        workspace.remove_document(random_uri)
 
     def references(self, doc_uri, position, exclude_declaration):
         return flatten(self._hook(
@@ -404,6 +476,16 @@ def m_text_document__did_close(self, textDocument=None, **_kwargs):
         workspace.publish_diagnostics(textDocument['uri'], [])
         workspace.rm_document(textDocument['uri'])
 
+    # TODO define params
+    def m_notebook_document__did_open(self, notebookDocument=None, cellTextDocuments=[], **_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'))
+        log.debug(f">>> cellTextDocuments: {cellTextDocuments}")
+        for cell in cellTextDocuments:
+            workspace.put_cell_document(cell['uri'], cell['languageId'], cell['text'], version=cell.get('version'))
+        # self._hook('pylsp_document_did_open', textDocument['uri'])  # This hook seems only relevant for rope
+        self.lint(notebookDocument['uri'], is_saved=True)
+
     def m_text_document__did_open(self, textDocument=None, **_kwargs):
         workspace = self._match_uri_to_workspace(textDocument['uri'])
         workspace.put_document(textDocument['uri'], textDocument['text'], version=textDocument.get('version'))

pylsp/workspace.py

@@ -105,12 +105,21 @@ def get_document(self, doc_uri):
         """
         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 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)
 
@@ -257,6 +266,29 @@ def _create_document(self, doc_uri, source=None, version=None):
             extra_sys_path=self.source_roots(path),
             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:
@@ -441,3 +473,27 @@ def sys_path(self, environment_path=None, env_vars=None):
         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)
+
+# We inherit from Document for now to get the same API. However, cell document differ from the text documents in that
+# as they have a language id.
+class Cell(Document):
+    """Represents a cell in a notebook."""
+
+    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

test/test_language_server.py

@@ -12,6 +12,7 @@
 import pytest
 
 from pylsp.python_lsp import start_io_lang_server, PythonLSPServer
+from pylsp.lsp import NotebookCellKind
 
 CALL_TIMEOUT = 10
 RUNNING_IN_CI = bool(os.environ.get('CI'))
@@ -118,3 +119,70 @@ def test_not_exit_without_check_parent_process_flag(client_server):  # pylint: d
 def test_missing_message(client_server):  # pylint: disable=redefined-outer-name
     with pytest.raises(JsonRpcMethodNotFound):
         client_server._endpoint.request('unknown_method').result(timeout=CALL_TIMEOUT)
+
+
+# TODO: make this assert on content of diagnostics message
+# Run this test if you want to see the diagnostics messages of an LSP server
+def test_text_document__did_open(client_server):
+    client_server._endpoint.request('initialize', {
+        'processId': 1234,
+        'rootPath': os.path.dirname(__file__),
+        'initializationOptions': {}
+    }).result(timeout=CALL_TIMEOUT)
+    client_server._endpoint.notify('textDocument/didOpen', {
+        'textDocument': {
+            'uri': os.path.join(os.path.dirname(__file__)),
+            'version': 1,
+            'text': 'import sys\nx=2\ny=x+2'
+        }
+    })
+
+# TODO: flesh the unit test out so that it tests the following:
+# The workspace is updated with the new notebook document and two cell documents
+def test_notebook_document__did_open(client_server):
+    client_server._endpoint.request('initialize', {
+        'processId': 1234,
+        'rootPath': os.path.dirname(__file__),
+        'initializationOptions': {}
+    }).result(timeout=CALL_TIMEOUT)
+    client_server._endpoint.notify('notebookDocument/didOpen', {
+        'notebookDocument': {
+            'uri': 'notebook_uri',
+            'notebookType': 'jupyter-notebook',
+            'cells': [
+                {
+                    'kind': NotebookCellKind.Code,
+                    'document': "cell_1_uri",
+                },
+                # TODO: add markdown cell support later
+                # {
+                #     'kind': NotebookCellKind.Markup,
+                #     'document': "cell_2_uri",
+                # },
+                {
+                    'kind': NotebookCellKind.Code,
+                    'document': "cell_3_uri",
+                }
+            ]
+        },
+        'cellTextDocuments': [
+            {
+                'uri': 'cell_1_uri',
+                'languageId': 'python',
+                'text': 'import sys',
+            },
+            # {
+            #     'uri': 'cell_2_uri',
+            #     'languageId': 'markdown',
+            #     'text': '# Title\n\n Some text',
+            # },
+            {
+                'uri': 'cell_3_uri',
+                'languageId': 'python',
+                'text': 'x = 2\ny = x + 2\nprint(z)',
+            }
+        ]
+    })
+#     time.sleep(0.1)
+#     # Assert that the documents are created in the workspace
+#     assert len(client_server.workspaces) == 1