handle EOF newline and an extra test

Author: masad-frost

pylsp/plugins/yapf_format.py

@@ -100,16 +100,18 @@ def _format(document, lines=None, options=None):
     diff = next(patch_generator)
     patch_generator.close()
 
-    # To keep things simple our text edits will be line based
-    # and uncompacted
+    # To keep things simple our text edits will be line based.
+    # We will also return the edits uncompacted, meaning a
+    # line replacement will come in as a line remove followed
+    # by a line add instead of a line replace.
     textEdits = []
     # keep track of line number since additions
     # don't include the line number it's being added
     # to in diffs. lsp is 0-indexed so we'll start with -1
     prev_line_no = -1
     for change in diff.changes:
         if change.old and change.new:
-            # no change
+            # old and new are the same line, no change
             # diffs are 1-indexed
             prev_line_no = change.old - 1
         elif change.new:
@@ -149,4 +151,23 @@ def _format(document, lines=None, options=None):
             })
             prev_line_no = lsp_line_no
 
+    # diffs don't include EOF newline https://github.com/google/yapf/issues/1008
+    # we'll add it ourselves if our document doesn't already have it and the diff
+    # does not change the last line.
+    if not source.endswith(eol_chars) and diff.changes \
+        and diff.changes[-1].old and diff.changes[-1].new:
+        textEdits.append({
+            'range': {
+                'start': {
+                    'line': prev_line_no,
+                    'character': 0
+                },
+                'end': {
+                    'line': prev_line_no + 1,
+                    'character': 0
+                }
+            },
+            'newText': diff.changes[-1].line + eol_chars
+        })
+
     return textEdits

test/plugins/test_yapf_format.py

@@ -68,29 +68,42 @@ def test_config_file(tmpdir, workspace):
 def test_line_endings(workspace, newline):
     doc = Document(DOC_URI, workspace, f'import os;import sys{2 * newline}dict(a=1)')
     res = pylsp_format_document(doc)
-
+    
     assert apply_text_edits(doc, res) == f'import os{newline}import sys{2 * newline}dict(a=1){newline}'
 
 
 def test_format_with_tab_size_option(workspace):
     doc = Document(DOC_URI, workspace, FOUR_SPACE_DOC)
     res = pylsp_format_document(doc, {"tabSize": "8"})
 
-    assert len(res) == 1
-    assert res[0]['newText'] == FOUR_SPACE_DOC.replace("    ", "        ")
+    assert apply_text_edits(doc, res) == FOUR_SPACE_DOC.replace("    ", "        ")
 
 
 def test_format_with_insert_spaces_option(workspace):
     doc = Document(DOC_URI, workspace, FOUR_SPACE_DOC)
     res = pylsp_format_document(doc, {"insertSpaces": False})
 
-    assert len(res) == 1
-    assert res[0]['newText'] == FOUR_SPACE_DOC.replace("    ", "\t")
+    assert apply_text_edits(doc, res) == FOUR_SPACE_DOC.replace("    ", "\t")
 
 
 def test_format_with_yapf_specific_option(workspace):
     doc = Document(DOC_URI, workspace, FOUR_SPACE_DOC)
     res = pylsp_format_document(doc, {"USE_TABS": True})
 
-    assert len(res) == 1
-    assert res[0]['newText'] == FOUR_SPACE_DOC.replace("    ", "\t")
+    assert apply_text_edits(doc, res) == FOUR_SPACE_DOC.replace("    ", "\t")
+
+def test_format_returns_text_edit_per_line(workspace):
+    single_space_indent = """def wow():
+ print("x")
+ print("hi")"""
+    doc = Document(DOC_URI, workspace, single_space_indent)
+    res = pylsp_format_document(doc)
+
+    # two removes and two adds
+    assert len(res) == 4
+    assert res[0]['newText'] == ""
+    assert res[1]['newText'] == ""
+    assert res[2]['newText'] == "    print(\"x\")\n"
+    assert res[3]['newText'] == "    print(\"hi\")\n"
+
+