Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add notebooks suppport to pylsp #389

Merged
merged 27 commits into from
Aug 9, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
70076fa
README.md
tkrabel-db Jun 5, 2023
ba0f9ad
Add notebook support
tkrabel-db Jun 7, 2023
d5bf4f9
Add example messages for reference
tkrabel-db Jun 7, 2023
78233c7
lsp_types: add DocumentUri and use correctly
tkrabel-db Jun 7, 2023
09d6879
add diagnostics support
tkrabel-db Jun 11, 2023
500e120
support editing, adding, and removing cells
tkrabel-db Jun 12, 2023
b407a31
add unit tests for notebook document messages
tkrabel-db Jun 13, 2023
53e4ad6
add test_notebook_document
tkrabel-db Jun 13, 2023
c891cf0
fix unit tests and pylint
tkrabel-db Jun 13, 2023
3dcaf9f
fix pytest test_notebook_document.py
tkrabel-db Jun 14, 2023
ab3bcee
Fix pylint issues:
tkrabel-db Jun 14, 2023
2dd7165
support notebookDocument__did_close
tkrabel-db Jun 14, 2023
7722239
Add notebookDocumentSync to capabilities
tkrabel-db Jun 15, 2023
6c00bfc
Add notebookDocumentSync to capabilities
tkrabel-db Jun 15, 2023
60517c1
fix: publishDiagnostics line starts at line 0
tkrabel-db Jul 10, 2023
ecc27ee
fix: publishDiagnostics starts at 0 and newlines are counted correctly
tkrabel-db Jul 10, 2023
25cdfab
fix: publishDiagnostics starts at 0 and newlines are counted correctly
tkrabel-db Jul 10, 2023
2b35c95
fix: publishDiagnostics starts at 0 and newlines are counted correctly
tkrabel-db Jul 10, 2023
69a1621
fix: publishDiagnostics starts at 0 and newlines are counted correctly
tkrabel-db Jul 10, 2023
1e9a51f
fix: publishDiagnostics starts at 0 and newlines are counted correctly
tkrabel-db Jul 10, 2023
4f140da
fix: publishDiagnostics starts at 0 and newlines are counted correctly
tkrabel-db Jul 10, 2023
e79bbb2
feat: close cell if deleted
tkrabel-db Jul 20, 2023
ce9f458
skip tests on windows as it's flaky on py3.7
tkrabel-db Jul 22, 2023
5434d1f
fix test_notebook_document__did_change: need to wait for 2 calls to d…
tkrabel-db Jul 24, 2023
ceb193d
Update test/data/publish_diagnostics_message_examples/example_1.json
tkrabel-db Jul 29, 2023
e7df839
remove aliases
tkrabel-db Jul 29, 2023
2dd05b7
address comments
tkrabel-db Jul 29, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 5 additions & 0 deletions pylsp/lsp.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,3 +90,8 @@ class TextDocumentSyncKind:
NONE = 0
FULL = 1
INCREMENTAL = 2


class NotebookCellKind:
Markup = 1
Code = 2
172 changes: 166 additions & 6 deletions pylsp/python_lsp.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
import os
import socketserver
import threading
import uuid
from typing import List, Dict, Any
import ujson as json

from pylsp_jsonrpc.dispatchers import MethodDispatcher
Expand All @@ -14,7 +16,7 @@

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__)
Expand Down Expand Up @@ -266,6 +268,11 @@ def capabilities(self):
},
'openClose': True,
},
'notebookDocumentSync': {
'notebookSelector': {
'cells': [{'language': 'python'}]
}
},
'workspace': {
'workspaceFolders': {
'supported': True,
Expand Down Expand Up @@ -375,11 +382,79 @@ 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): # 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)
tkrabel-db marked this conversation as resolved.
Show resolved Hide resolved

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(
Expand All @@ -399,6 +474,91 @@ def folding(self, doc_uri):
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'], [])
Expand Down
87 changes: 86 additions & 1 deletion pylsp/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
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
Expand All @@ -35,6 +35,8 @@ def wrapper(self, *args, **kwargs):

class Workspace:

# pylint: disable=too-many-public-methods

M_PUBLISH_DIAGNOSTICS = 'textDocument/publishDiagnostics'
M_PROGRESS = '$/progress'
M_INITIALIZE_PROGRESS = 'window/workDoneProgress/create'
Expand Down Expand Up @@ -105,12 +107,30 @@ 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 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)

Expand Down Expand Up @@ -258,6 +278,29 @@ def _create_document(self, doc_uri, source=None, version=None):
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()
Expand Down Expand Up @@ -441,3 +484,45 @@ 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)

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'))
Loading
Loading