8000 Add go to definition support for notebook cell documents · python-lsp/python-lsp-server@7f496bb · GitHub
[go: up one dir, main page]
Skip to content

Commit 7f496bb

Browse files
jasongroutjasongrout
committed
Add go to definition support for notebook cell documents
1 parent 2757c3a commit 7f496bb

File tree

3 files changed

+162
-7
lines changed

3 files changed

+162
-7
lines changed

‎pylsp/python_lsp.py

Lines changed: 77 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
from . import lsp, _utils, uris
1818
from .config import config
19-
from .workspace import Workspace, Document, Notebook
19+
from .workspace import Workspace, Document, Notebook, Cell
2020
from ._version import __version__
2121

2222
log = logging.getLogger(__name__)
@@ -484,7 +484,7 @@ def m_notebook_document__did_open(self, notebookDocument=None, cellTextDocuments
484484
metadata=notebookDocument.get('metadata')
485485
)
486486
for cell in (cellTextDocuments or []):
487-
workspace.put_cell_document(cell['uri'], cell['languageId'], cell['text'], version=cell.get('version'))
487+
workspace.put_cell_document(cell['uri'], notebookDocument['uri'], cell['languageId'], cell['text'], version=cell.get('version'))
488488
self.lint(notebookDocument['uri'], is_saved=True)
489489

490490
def m_notebook_document__did_close(self, notebookDocument=None, cellTextDocuments=None, **_kwargs):
@@ -526,7 +526,8 @@ def m_notebook_document__did_change(self, notebookDocument=None, change=None, **
526526
# Cell documents
527527
for cell_document in structure['didOpen']:
528528
workspace.put_cell_document(
529-
cell_document['uri'],
529+
cell_document['uri'],
530+
notebookDocument['uri'],
530531
cell_document['languageId'],
531532
cell_document['text'],
532533
cell_document.get('version')
@@ -593,7 +594,80 @@ def m_text_document__code_lens(self, textDocument=None, **_kwargs):
593594
def m_text_document__completion(self, textDocument=None, position=None, **_kwargs):
594595
return self.completions(textDocument['uri'], position)
595596

597+
def _cell_document__definition(self, cellDocument=None, position=None, **_kwargs):
598+
# First, we create a temp TextDocument to send to the hook that represents the whole notebook
599+
# contents.
600+
workspace = self._match_uri_to_workspace(cellDocument.notebook_uri)
601+
notebookDocument = workspace.get_maybe_document(cellDocument.notebook_uri)
602+
if notebookDocument is None:
603+
raise ValueError("Invalid notebook document")
604+
605+
random_uri = str(uuid.uuid4())
606+
# cell_list helps us map the diagnostics back to the correct cell later.
607+
cell_list: List[Dict[str, Any]] = []
608+
609+
offset = 0
610+
total_source = ""
611+
for cell in notebookDocument.cells:
612+
cell_uri = cell['document']
613+
cell_document = workspace.get_cell_document(cell_uri)
614+
615+
num_lines = cell_document.line_count
616+
617+
data = {
618+
'uri': cell_uri,
619+
'line_start': offset,
620+
'line_end': offset + num_lines - 1,
621+
'source': cell_document.source
622+
}
623+
624+
if position is not None and cell_uri == cellDocument.uri:
625+
position['line'] += offset
626+
627+
cell_list.append(data)
628+
if offset == 0:
629+
total_source = cell_document.source
630+
else:
631+
total_source += ("\n" + cell_document.source)
632+
633+
offset += num_lines
634+
635+
# TODO: make a workspace temp document context manager that yields the random uri and cleans up afterwards
636+
workspace.put_document(random_uri, total_source)
637+
log.info(f'Making new document {random_uri}')
638+
try:
639+
definitions = self.definitions(random_uri, position)
640+
log.info(f'Got definitions: {definitions}')
641+
642+
# {
643+
# 'uri': uris.uri_with(document.uri, path=str(d.module_path)),
644+
# 'range': {
645+
# 'start': {'line': d.line - 1, 'character': d.column},
646+
# 'end': {'line': d.line - 1, 'character': d.column + len(d.name)},
647+
# }
648+
# }
649+
print(definitions)
650+
for definition in definitions:
651+
# TODO: a better test for if a definition is the random_uri
652+
if random_uri in definition['uri']:
653+
# Find the cell the start is in
654+
for cell in cell_list:
655+
# TODO: perhaps it is more correct to check definition['range']['end']['line'] <= cell['line_end'], but
656+
# that would mess up if a definition was split over cells
657+
if cell['line_start'] <= definition['range']['start']['line'] <= cell['line_end']:
658+
definition['uri'] = cell['uri']
659+
definition['range']['start']['line'] -= cell['line_start']
660+
definition['range']['end']['line'] -= cell['line_start']
661+
return definitions
662+
finally:
663+
workspace.rm_document(random_uri)
664+
596665
def m_text_document__definition(self, textDocument=None, position=None, **_kwargs):
666+
# textDocument here is just a dict with a uri
667+
workspace = self._match_uri_to_workspace(textDocument['uri'])
668+
document = workspace.get_document(textDocument['uri'])
669+
if isinstance(document, Cell):
670+
return self._cell_document__definition(document, position, **_kwargs)
597671
return self.definitions(textDocument['uri'], position)
598672

599673
def m_text_document__document_highlight(self, textDocument=None, position=None, **_kwargs):

‎pylsp/workspace.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -128,8 +128,8 @@ def remove_notebook_cells(self, doc_uri, start, delete_count):
128128
def update_notebook_metadata(self, doc_uri, metadata):
129129
self._docs[doc_uri].metadata = metadata
130130

131-
def put_cell_document(self, doc_uri, language_id, source, version=None):
132-
self._docs[doc_uri] = self._create_cell_document(doc_uri, language_id, source, version)
131+
def put_cell_document(self, doc_uri, notebook_uri, language_id, source, version=None):
132+
self._docs[doc_uri] = self._create_cell_document(doc_uri, notebook_uri, language_id, source, version)
133133

134134
def rm_document(self, doc_uri):
135135
self._docs.pop(doc_uri)
@@ -305,11 +305,12 @@ def _create_notebook_document(self, doc_uri, notebook_type, cells, version=None,
305305
metadata=metadata
306306
)
307307

308-
def _create_cell_document(self, doc_uri, language_id, source=None, version=None):
308+
def _create_cell_document(self, doc_uri, notebook_uri, language_id, source=None, version=None):
309309
# TODO: remove what is unnecessary here.
310310
path = uris.to_fs_path(doc_uri)
311311
return Cell(
312312
doc_uri,
313+
notebook_uri=notebook_uri,
313314
language_id=language_id,
314315
workspace=self,
315316
source=source,
@@ -538,10 +539,11 @@ class Cell(Document):
538539
they have a language id.
539540
"""
540541

541-
def __init__(self, uri, language_id, workspace, source=None, version=None, local=True, extra_sys_path=None,
542+
def __init__(self, uri, notebook_uri, language_id, workspace, source=None, version=None, local=True, extra_sys_path=None,
542543
rope_project_builder=None):
543544
super().__init__(uri, workspace, source, version, local, extra_sys_path, rope_project_builder)
544545
self.language_id = language_id
546+
self.notebook_uri = notebook_uri
545547

546548
@property
547549
@lock

‎test/test_notebook_document.py

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -582,3 +582,82 @@ def test_notebook__did_close(
582582
)
583583
wait_for_condition(lambda: mock_notify.call_count >= 2)
584584
assert len(server.workspace.documents) == 0
585+
586+
def test_notebook_definition(
587+
client_server_pair,
588+
): # pylint: disable=redefined-outer-name
589+
client, server = client_server_pair
590+
client._endpoint.request(
591+
"initialize",
592+
{
593+
"processId": 1234,
594+
"rootPath": os.path.dirname(__file__),
595+
"initializationOptions": {},
596+
},
597+
).result(timeout=CALL_TIMEOUT_IN_SECONDS)
598+
599+
# Open notebook
600+
with patch.object(server._endpoint, "notify") as mock_notify:
601+
client._endpoint.notify(
602+
"notebookDocument/didOpen",
603+
{
604+
"notebookDocument": {
605+
"uri": "notebook_uri",
606+
"notebookType": "jupyter-notebook",
607+
"cells": [
608+
{
609+
"kind": NotebookCellKind.Code,
610+
"document": "cell_1_uri",
611+
},
612+
{
613+
"kind": NotebookCellKind.Code,
614+
"document": "cell_2_uri",
615+
},
616+
],
617+
},
618+
"cellTextDocuments": [
619+
{
620+
"uri": "cell_1_uri",
621+
"languageId": "python",
622+
"text": "y=2\nx=1",
623+
},
624+
{
625+
"uri": "cell_2_uri",
626+
"languageId": "python",
627+
"text": "x",
628+
},
629+
],
630+
},
631+
)
632+
# wait for expected diagnostics messages
633+
wait_for_condition(lambda: mock_notify.call_count >= 2)
634+
assert len(server.workspace.documents) == 3
635+
for uri in ["cell_1_uri", "cell_2_uri", "notebook_uri"]:
636+
assert uri in server.workspace.documents
637+
638+
future = client._endpoint.request(
639+
"textDocument/definition",
640+
{
641+
"textDocument": {
642+
"uri": "cell_2_uri",
643+
},
644+
"position": {
645+
"line": 0,
646+
"character": 1
647+
}
648+
},
649+
)
650+
result = future.result(CALL_TIMEOUT_IN_SECONDS)
651+
assert result == [{
652+
'uri': 'cell_1_uri',
653+
'range': {
654+
'start': {
655+
'line': 1,
656+
'character': 0
657+
},
658+
'end': {
659+
'line': 1,
660+
'character': 1
661+
}
662+
}
663+
}]

0 commit comments

Comments
 (0)
0