Skip to content

Commit ecc27ee

Browse files
tkrabel-dbtkrabel-db
committed
fix: publishDiagnostics starts at 0 and newlines are counted correctly
1 parent 60517c1 commit ecc27ee

File tree

3 files changed

+162
-35
lines changed

3 files changed

+162
-35
lines changed

‎pylsp/python_lsp.py

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -416,34 +416,38 @@ def _lint_notebook_document(self, notebook_document, workspace):
416416
cell_uri = cell['document']
417417
cell_document = workspace.get_cell_document(cell_uri)
418418

419-
num_lines = len(cell_document.lines)
419+
num_lines = cell_document.line_count
420420

421421
data = {
422422
'uri': cell_uri,
423-
'line_start': offset + 1,
424-
'line_end': offset + num_lines,
423+
'line_start': offset,
424+
'line_end': offset + num_lines - 1,
425425
'source': cell_document.source
426426
}
427427

428428
cell_list.append(data)
429-
total_source = total_source + "\n" + cell_document.source
429+
if offset == 0:
430+
total_source = cell_document.source
431+
else:
432+
maybe_newline = "" if total_source.endswith("\n") else "\n"
433+
total_source += (maybe_newline + cell_document.source)
430434

431435
offset += num_lines
432436

433437
workspace.put_document(random_uri, total_source)
434438
document_diagnostics = flatten(self._hook('pylsp_lint', random_uri, is_saved=True))
435439

436-
# Now we need to map the diagnostics back to the correct cell and
437-
# publish them.
440+
# Now we need to map the diagnostics back to the correct cell and publish them.
441+
# Note: this is O(n*m) in the number of cells and diagnostics, respectively.
438442
for cell in cell_list:
439443
cell_diagnostics = []
440444
for diagnostic in document_diagnostics:
441-
if diagnostic['range']['start']['line'] > cell['line_end']:
442-
break
445+
if diagnostic['range']['start']['line'] > cell['line_end'] \
446+
or diagnostic['range']['end']['line'] < cell['line_start']:
447+
continue
443448
diagnostic['range']['start']['line'] = diagnostic['range']['start']['line'] - cell['line_start']
444449
diagnostic['range']['end']['line'] = diagnostic['range']['end']['line'] - cell['line_start']
445450
cell_diagnostics.append(diagnostic)
446-
document_diagnostics.pop(0)
447451

448452
workspace.publish_diagnostics(cell['uri'], cell_diagnostics)
449453

‎pylsp/workspace.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -515,3 +515,23 @@ def __init__(self, uri, language_id, workspace, source=None, version=None, local
515515
rope_project_builder=None):
516516
super().__init__(uri, workspace, source, version, local, extra_sys_path, rope_project_builder)
517517
self.language_id = language_id
518+
519+
@property
520+
@lock
521+
def line_count(self):
522+
""""
523+
Return the number of lines in the cell document.
524+
525+
This function is used to combine all cell documents into one text document. Note that we need to count a
526+
newline at the end of a non-empty text line as an extra line.
527+
528+
Examples:
529+
- "x\n" is a cell with two lines, "x" and ""
530+
- "x" is a cell with one line, "x"
531+
- "\n" is a cell with one line, ""
532+
"""
533+
if len(self.lines) == 0:
534+
return 1
535+
536+
last_line = self.lines[-1]
537+
return len(self.lines) + (last_line != "\n" and last_line.endswith("\n"))

‎test/test_notebook_document.py

Lines changed: 129 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@
1010
from pylsp.python_lsp import PythonLSPServer
1111
from pylsp.lsp import NotebookCellKind
1212

13-
CALL_TIMEOUT = 10
13+
CALL_TIMEOUT_IN_SECONDS = 10
1414

1515

16-
def wait_for_condition(condition, timeout=CALL_TIMEOUT):
16+
def wait_for_condition(condition, timeout=CALL_TIMEOUT_IN_SECONDS):
1717
"""Wait for a condition to be true, or timeout."""
1818
start_time = time.time()
1919
while not condition():
@@ -51,9 +51,9 @@ def client_server_pair():
5151

5252
yield (client_server_pair_obj.client, client_server_pair_obj.server)
5353

54-
shutdown_response = client_server_pair_obj.client._endpoint.request("shutdown").result(
55-
timeout=CALL_TIMEOUT
56-
)
54+
shutdown_response = client_server_pair_obj.client._endpoint.request(
55+
"shutdown"
56+
).result(timeout=CALL_TIMEOUT_IN_SECONDS)
5757
assert shutdown_response is None
5858
client_server_pair_obj.client._endpoint.notify("exit")
5959

@@ -67,13 +67,15 @@ def test_initialize(client_server_pair): # pylint: disable=redefined-outer-name
6767
"rootPath": os.path.dirname(__file__),
6868
"initializationOptions": {},
6969
},
70-
).result(timeout=CALL_TIMEOUT)
70+
).result(timeout=CALL_TIMEOUT_IN_SECONDS)
7171
assert server.workspace is not None
7272
assert "capabilities" in response
7373
# TODO: assert that notebook capabilities are in response
7474

7575

76-
def test_notebook_document__did_open(client_server_pair): # pylint: disable=redefined-outer-name
76+
def test_notebook_document__did_open(
77+
client_server_pair,
78+
): # pylint: disable=redefined-outer-name
7779
client, server = client_server_pair
7880
client._endpoint.request(
7981
"initialize",
@@ -82,7 +84,7 @@ def test_notebook_document__did_open(client_server_pair): # pylint: disable=red
8284
"rootPath": os.path.dirname(__file__),
8385
"initializationOptions": {},
8486
},
85-
).result(timeout=CALL_TIMEOUT)
87+
).result(timeout=CALL_TIMEOUT_IN_SECONDS)
8688

8789
with patch.object(server._endpoint, "notify") as mock_notify:
8890
client._endpoint.notify(
@@ -100,28 +102,70 @@ def test_notebook_document__did_open(client_server_pair): # pylint: disable=red
100102
"kind": NotebookCellKind.Code,
101103
"document": "cell_2_uri",
102104
},
105+
{
106+
"kind": NotebookCellKind.Code,
107+
"document": "cell_3_uri",
108+
},
109+
{
110+
"kind": NotebookCellKind.Code,
111+
"document": "cell_4_uri",
112+
},
113+
{
114+
"kind": NotebookCellKind.Code,
115+
"document": "cell_5_uri",
116+
},
103117
],
104118
},
119+
# Test as many edge cases as possible for the diagnostics message
105120
"cellTextDocuments": [
106121
{
107122
"uri": "cell_1_uri",
108123
"languageId": "python",
109-
"text": "import sys",
124+
"text": "",
110125
},
111126
{
112127
"uri": "cell_2_uri",
113128
"languageId": "python",
114-
"text": "x = 1",
129+
"text": "\n",
130+
},
131+
{
132+
"uri": "cell_3_uri",
133+
"languageId": "python",
134+
"text": "\nimport sys\n\nabc\n\n",
135+
},
136+
{
137+
"uri": "cell_4_uri",
138+
"languageId": "python",
139+
"text": "x",
140+
},
141+
{
142+
"uri": "cell_5_uri",
143+
"languageId": "python",
144+
"text": "y\n",
115145
},
116146
],
117147
},
118148
)
119-
wait_for_condition(lambda: mock_notify.call_count >= 2)
149+
wait_for_condition(lambda: mock_notify.call_count >= 5)
120150
expected_call_args = [
121151
call(
122152
"textDocument/publishDiagnostics",
123153
params={
124154
"uri": "cell_1_uri",
155+
"diagnostics": [],
156+
},
157+
),
158+
call(
159+
"textDocument/publishDiagnostics",
160+
params={
161+
"uri": "cell_2_uri",
162+
"diagnostics": [],
163+
},
164+
),
165+
call(
166+
"textDocument/publishDiagnostics",
167+
params={
168+
"uri": "cell_3_uri",
125169
"diagnostics": [
126170
{
127171
"source": "pyflakes",
@@ -131,33 +175,70 @@ def test_notebook_document__did_open(client_server_pair): # pylint: disable=red
131175
},
132176
"message": "'sys' imported but unused",
133177
"severity": 2,
134-
}
178+
},
179+
{
180+
"source": "pyflakes",
181+
"range": {
182+
"start": {"line": 3, "character": 0},
183+
"end": {"line": 3, "character": 4},
184+
},
185+
"message": "undefined name 'abc'",
186+
"severity": 1,
187+
},
188+
{
189+
"source": "pycodestyle",
190+
"range": {
191+
"start": {"line": 1, "character": 0},
192+
"end": {"line": 1, "character": 11},
193+
},
194+
"message": "E303 too many blank lines (3)",
195+
"code": "E303",
196+
"severity": 2,
197+
},
135198
],
136199
},
137200
),
138201
call(
139202
"textDocument/publishDiagnostics",
140203
params={
141-
"uri": "cell_2_uri",
204+
"uri": "cell_4_uri",
142205
"diagnostics": [
143206
{
144-
"source": "pycodestyle",
207+
"source": "pyflakes",
145208
"range": {
146-
"start": {"line": 0, "character": 5},
147-
"end": {"line": 0, "character": 5},
209+
"start": {"line": 0, "character": 0},
210+
"end": {"line": 0, "character": 2},
148211
},
149-
"message": "W292 no newline at end of file",
150-
"code": "W292",
151-
"severity": 2,
152-
}
212+
"message": "undefined name 'x'",
213+
"severity": 1,
214+
},
215+
],
216+
},
217+
),
218+
call(
219+
"textDocument/publishDiagnostics",
220+
params={
221+
"uri": "cell_5_uri",
222+
"diagnostics": [
223+
{
224+
"source": "pyflakes",
225+
"range": {
226+
"start": {"line": 0, "character": 0},
227+
"end": {"line": 0, "character": 2},
228+
},
229+
"message": "undefined name 'y'",
230+
"severity": 1,
231+
},
153232
],
154233
},
155234
),
156235
]
157236
mock_notify.assert_has_calls(expected_call_args)
158237

159238

160-
def test_notebook_document__did_change(client_server_pair): # pylint: disable=redefined-outer-name
239+
def test_notebook_document__did_change(
240+
client_server_pair,
241+
): # pylint: disable=redefined-outer-name
161242
client, server = client_server_pair
162243
client._endpoint.request(
163244
"initialize",
@@ -166,7 +247,7 @@ def test_notebook_document__did_change(client_server_pair): # pylint: disable=r
166247
"rootPath": os.path.dirname(__file__),
167248
"initializationOptions": {},
168249
},
169-
).result(timeout=CALL_TIMEOUT)
250+
).result(timeout=CALL_TIMEOUT_IN_SECONDS)
170251

171252
# Open notebook
172253
with patch.object(server._endpoint, "notify") as mock_notify:
@@ -274,7 +355,17 @@ def test_notebook_document__did_change(client_server_pair): # pylint: disable=r
274355
},
275356
"message": "'sys' imported but unused",
276357
"severity": 2,
277-
}
358+
},
359+
{
360+
"source": "pycodestyle",
361+
"range": {
362+
"start": {"line": 0, "character": 10},
363+
"end": {"line": 0, "character": 10},
364+
},
365+
"message": "W292 no newline at end of file",
366+
"code": "W292",
367+
"severity": 2,
368+
},
278369
],
279370
},
280371
)
@@ -349,7 +440,17 @@ def test_notebook_document__did_change(client_server_pair): # pylint: disable=r
349440
},
350441
"message": "undefined name 'x'",
351442
"severity": 1,
352-
}
443+
},
444+
{
445+
"source": "pycodestyle",
446+
"range": {
447+
"start": {"line": 0, "character": 1},
448+
"end": {"line": 0, "character": 1},
449+
},
450+
"message": "W292 no newline at end of file",
451+
"code": "W292",
452+
"severity": 2,
453+
},
353454
],
354455
},
355456
),
@@ -406,7 +507,9 @@ def test_notebook_document__did_change(client_server_pair): # pylint: disable=r
406507
mock_notify.assert_has_calls(expected_call_args)
407508

408509

409-
def test_notebook__did_close(client_server_pair): # pylint: disable=redefined-outer-name
510+
def test_notebook__did_close(
511+
client_server_pair,
512+
): # pylint: disable=redefined-outer-name
410513
client, server = client_server_pair
411514
client._endpoint.request(
412515
"initialize",
@@ -415,7 +518,7 @@ def test_notebook__did_close(client_server_pair): # pylint: disable=redefined-
415518
"rootPath": os.path.dirname(__file__),
416519
"initializationOptions": {},
417520
},
418-
).result(timeout=CALL_TIMEOUT)
521+
).result(timeout=CALL_TIMEOUT_IN_SECONDS)
419522

420523
# Open notebook
421524
with patch.object(server._endpoint, "notify") as mock_notify:

0 commit comments

Comments
 (0)