Skip to content

Commit 9244d5a

Browse files
seeMseeM
committed
wip trying decorator idea
1 parent 4ed862f commit 9244d5a

File tree

2 files changed

+115
-96
lines changed

2 files changed

+115
-96
lines changed

‎jedi_language_server/notebook_utils.py

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,44 @@
11
"""Utility functions for handling notebook documents."""
22

3+
from __future__ import annotations
4+
35
from collections import defaultdict
46
from typing import (
7+
TYPE_CHECKING,
8+
Callable,
59
Dict,
610
Iterable,
711
List,
812
NamedTuple,
913
Optional,
14+
Tuple,
15+
TypeVar,
1016
Union,
1117
cast,
1218
)
1319

1420
import attrs
1521
from lsprotocol.types import (
1622
AnnotatedTextEdit,
23+
CodeActionParams,
24+
CompletionParams,
25+
Hover,
1726
Location,
1827
NotebookDocument,
1928
OptionalVersionedTextDocumentIdentifier,
2029
Position,
2130
Range,
31+
RenameParams,
2232
TextDocumentEdit,
33+
TextDocumentPositionParams,
2334
TextEdit,
2435
)
2536
from pygls.uris import to_fs_path
2637
from pygls.workspace import TextDocument, Workspace
2738

39+
if TYPE_CHECKING:
40+
from .server import JediLanguageServer
41+
2842

2943
def notebook_coordinate_mapper(
3044
workspace: Workspace,
@@ -239,3 +253,85 @@ def cell_index(workspace: Workspace, cell_uri: str) -> int:
239253
index = notebook.cell_index(cell_uri)
240254
assert index is not None
241255
return index
256+
257+
258+
NotebookSupportedParams = Union[
259+
CodeActionParams,
260+
CompletionParams,
261+
RenameParams,
262+
TextDocumentPositionParams,
263+
]
264+
T_params = TypeVar(
265+
"T_params",
266+
bound=NotebookSupportedParams,
267+
)
268+
269+
270+
def notebook_text_document_and_params(
271+
workspace: Workspace,
272+
params: T_params,
273+
) -> Tuple[TextDocument, T_params]:
274+
notebook = notebook_coordinate_mapper(
275+
workspace, cell_uri=params.text_document.uri
276+
)
277+
if notebook is None:
278+
raise ValueError(
279+
f"Notebook not found with cell URI: {params.text_document.uri}"
280+
)
281+
document = TextDocument(uri=notebook._document.uri, source=notebook.source)
282+
283+
position = getattr(params, "position", None)
284+
if position is not None:
285+
notebook_position = notebook.notebook_position(
286+
params.text_document.uri, position
287+
)
288+
params = attrs.evolve(params, position=notebook_position) # type: ignore[arg-type]
289+
290+
range = getattr(params, "range", None)
291+
if range is not None:
292+
notebook_range = notebook.notebook_range(
293+
params.text_document.uri, range
294+
)
295+
params = attrs.evolve(params, range=notebook_range) # type: ignore[arg-type]
296+
297+
return document, params
298+
299+
300+
T = TypeVar("T")
301+
302+
303+
def with_notebook_support(
304+
f: Callable[
305+
[JediLanguageServer, T_params, Optional[TextDocument]],
306+
T,
307+
],
308+
) -> Callable[[JediLanguageServer, T_params], T]:
309+
def wrapped(server: JediLanguageServer, params: T_params) -> T:
310+
document, params = notebook_text_document_and_params(
311+
server.workspace, params
312+
)
313+
result = f(server, params, document)
314+
315+
if (
316+
isinstance(result, list)
317+
and result
318+
and isinstance(result[0], Location)
319+
):
320+
return cast(
321+
T, text_document_or_cell_locations(server.workspace, result)
322+
)
323+
324+
if isinstance(result, Hover) and result.range is not None:
325+
notebook_mapper = notebook_coordinate_mapper(
326+
server.workspace, cell_uri=params.text_document.uri
327+
)
328+
if notebook_mapper is None:
329+
return cast(T, result)
330+
location = notebook_mapper.cell_range(result.range)
331+
if location is None or location.uri != params.text_document.uri:
332+
return cast(T, result)
333+
return cast(T, attrs.evolve(result, range=location.range))
334+
335+
return result
336+
337+
return wrapped

‎jedi_language_server/server_ipynb.py

Lines changed: 19 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
1-
from typing import List, Optional, Tuple, TypeVar, Union
1+
from typing import List, Optional
22

3-
import attrs
43
from lsprotocol.types import (
54
TEXT_DOCUMENT_CODE_ACTION,
65
TEXT_DOCUMENT_COMPLETION,
@@ -26,136 +25,90 @@
2625
WorkspaceEdit,
2726
)
2827
from pygls.workspace import Workspace
29-
from pygls.workspace.text_document import TextDocument
3028

3129
from . import notebook_utils, server_utils
3230
from .server import SERVER, JediLanguageServer
3331

34-
T_params = TypeVar(
35-
"T_params",
36-
bound=Union[
37-
CodeActionParams,
38-
CompletionParams,
39-
RenameParams,
40-
TextDocumentPositionParams,
41-
],
42-
)
43-
4432

4533
@SERVER.notebook_feature(TEXT_DOCUMENT_COMPLETION)
4634
def completion(
4735
server: JediLanguageServer, params: CompletionParams
4836
) -> Optional[CompletionList]:
49-
document, params = _notebook_text_document_and_params(
50-
server.workspace, params
37+
return notebook_utils.with_notebook_support(server_utils.completion)(
38+
server, params
5139
)
52-
return server_utils.completion(server, params, document)
5340

5441

5542
@SERVER.notebook_feature(TEXT_DOCUMENT_SIGNATURE_HELP)
5643
def signature_help(
5744
server: JediLanguageServer, params: TextDocumentPositionParams
5845
) -> Optional[SignatureHelp]:
59-
document, params = _notebook_text_document_and_params(
60-
server.workspace, params
46+
return notebook_utils.with_notebook_support(server_utils.signature_help)(
47+
server, params
6148
)
62-
return server_utils.signature_help(server, params, document)
6349

6450

6551
@SERVER.notebook_feature(TEXT_DOCUMENT_DECLARATION)
6652
def declaration(
6753
server: JediLanguageServer, params: TextDocumentPositionParams
6854
) -> Optional[List[Location]]:
69-
document, params = _notebook_text_document_and_params(
70-
server.workspace, params
71-
)
72-
locations = server_utils.declaration(server, params, document)
73-
return notebook_utils.text_document_or_cell_locations(
74-
server.workspace, locations
55+
return notebook_utils.with_notebook_support(server_utils.declaration)(
56+
server, params
7557
)
7658

7759

7860
@SERVER.notebook_feature(TEXT_DOCUMENT_DEFINITION)
7961
def definition(
8062
server: JediLanguageServer, params: TextDocumentPositionParams
8163
) -> Optional[List[Location]]:
82-
document, params = _notebook_text_document_and_params(
83-
server.workspace, params
84-
)
85-
locations = server_utils.definition(server, params, document)
86-
return notebook_utils.text_document_or_cell_locations(
87-
server.workspace, locations
64+
return notebook_utils.with_notebook_support(server_utils.definition)(
65+
server, params
8866
)
8967

9068

9169
@SERVER.notebook_feature(TEXT_DOCUMENT_TYPE_DEFINITION)
9270
def type_definition(
9371
server: JediLanguageServer, params: TextDocumentPositionParams
9472
) -> Optional[List[Location]]:
95-
document, params = _notebook_text_document_and_params(
96-
server.workspace, params
97-
)
98-
locations = server_utils.type_definition(server, params, document)
99-
return notebook_utils.text_document_or_cell_locations(
100-
server.workspace, locations
73+
return notebook_utils.with_notebook_support(server_utils.type_definition)(
74+
server, params
10175
)
10276

10377

10478
def hover(
10579
server: JediLanguageServer, params: TextDocumentPositionParams
10680
) -> Optional[Hover]:
107-
document, params = _notebook_text_document_and_params(
108-
server.workspace, params
81+
return notebook_utils.with_notebook_support(server_utils.hover)(
82+
server, params
10983
)
110-
# TODO: Can we clean this up?
111-
# Maybe have a single notebook_to_cell and vice versa function that
112-
# can handle most needed lsprotocol types?
113-
hover = server_utils.hover(server, params, document)
114-
if hover is None or hover.range is None:
115-
return hover
116-
notebook_mapper = notebook_utils.notebook_coordinate_mapper(
117-
server.workspace, cell_uri=params.text_document.uri
118-
)
119-
if notebook_mapper is None:
120-
return hover
121-
location = notebook_mapper.cell_range(hover.range)
122-
if location is None or location.uri != params.text_document.uri:
123-
return hover
124-
return attrs.evolve(hover, range=location.range)
12584

12685

12786
@SERVER.notebook_feature(TEXT_DOCUMENT_REFERENCES)
12887
def references(
12988
server: JediLanguageServer, params: TextDocumentPositionParams
13089
) -> Optional[List[Location]]:
131-
document, params = _notebook_text_document_and_params(
132-
server.workspace, params
133-
)
134-
locations = server_utils.references(server, params, document)
135-
return notebook_utils.text_document_or_cell_locations(
136-
server.workspace, locations
90+
return notebook_utils.with_notebook_support(server_utils.references)(
91+
server, params
13792
)
13893

13994

14095
@SERVER.notebook_feature(TEXT_DOCUMENT_RENAME)
14196
def rename(
14297
server: JediLanguageServer, params: RenameParams
14398
) -> Optional[WorkspaceEdit]:
144-
document, params = _notebook_text_document_and_params(
145-
server.workspace, params
99+
return notebook_utils.with_notebook_support(server_utils.rename)(
100+
server, params
146101
)
147-
return server_utils.rename(server, params, document)
148102

149103

150104
@SERVER.notebook_feature(TEXT_DOCUMENT_CODE_ACTION)
151105
def code_action(
152106
server: JediLanguageServer,
153107
params: CodeActionParams,
154108
) -> Optional[List[CodeAction]]:
155-
document, params = _notebook_text_document_and_params(
156-
server.workspace, params
109+
return notebook_utils.with_notebook_support(server_utils.code_action)(
110+
server, params
157111
)
158-
return server_utils.code_action(server, params, document)
159112

160113

161114
# NOTEBOOK_DOCUMENT_DID_SAVE
@@ -244,36 +197,6 @@ def did_close_notebook_document_default(
244197
"""Actions run on notebookDocument/didClose: default."""
245198

246199

247-
def _notebook_text_document_and_params(
248-
workspace: Workspace,
249-
params: T_params,
250-
) -> Tuple[TextDocument, T_params]:
251-
notebook = notebook_utils.notebook_coordinate_mapper(
252-
workspace, cell_uri=params.text_document.uri
253-
)
254-
if notebook is None:
255-
raise ValueError(
256-
f"Notebook not found with cell URI: {params.text_document.uri}"
257-
)
258-
document = TextDocument(uri=notebook._document.uri, source=notebook.source)
259-
260-
position = getattr(params, "position", None)
261-
if position is not None:
262-
notebook_position = notebook.notebook_position(
263-
params.text_document.uri, position
264-
)
265-
params = attrs.evolve(params, position=notebook_position) # type: ignore[arg-type]
266-
267-
range = getattr(params, "range", None)
268-
if range is not None:
269-
notebook_range = notebook.notebook_range(
270-
params.text_document.uri, range
271-
)
272-
params = attrs.evolve(params, range=notebook_range) # type: ignore[arg-type]
273-
274-
return document, params
275-
276-
277200
def _cell_filename(
278201
workspace: Workspace,
279202
cell_uri: str,

0 commit comments

Comments
 (0)