Sendback diffs instead of the whole document with yapf formatting

Author: masad-frost

pylsp/plugins/yapf_format.py

@@ -7,6 +7,8 @@
 from yapf.yapflib import file_resources
 from yapf.yapflib.yapf_api import FormatCode
 
+import whatthepatch
+
 from pylsp import hookimpl
 from pylsp._utils import get_eol_chars
 
@@ -39,17 +41,16 @@ def pylsp_format_range(document, range):  # pylint: disable=redefined-builtin
 def _format(document, lines=None):
     # Yapf doesn't work with CR line endings, so we replace them by '\n'
     # and restore them below.
-    replace_cr = False
     source = document.source
     eol_chars = get_eol_chars(source)
     if eol_chars == '\r':
-        replace_cr = True
         source = source.replace('\r', '\n')
 
-    new_source, changed = FormatCode(
+    diff_txt, changed = FormatCode(
         source,
         lines=lines,
         filename=document.filename,
+        print_diff=True,
         style_config=file_resources.GetDefaultStyleForDir(
             os.path.dirname(document.path)
         )
@@ -58,16 +59,57 @@ def _format(document, lines=None):
     if not changed:
         return []
 
-    if replace_cr:
-        new_source = new_source.replace('\n', '\r')
-
-    # I'm too lazy at the moment to parse diffs into TextEdit items
-    # So let's just return the entire file...
-    return [{
-        'range': {
-            'start': {'line': 0, 'character': 0},
-            # End char 0 of the line after our document
-            'end': {'line': len(document.lines), 'character': 0}
-        },
-        'newText': new_source
-    }]
+    patch_generator = whatthepatch.parse_patch(diff_txt)
+    diff = next(patch_generator)
+    patch_generator.close()
+
+    # To keep things simple our text edits will be line based
+    # and uncompacted
+    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
+            # diffs are 1-indexed
+            prev_line_no = change.old - 1
+        elif change.new:
+            # addition
+            textEdits.append({
+                'range': {
+                    'start': {
+                        'line': prev_line_no + 1,
+                        'character': 0
+                    },
+                    'end': {
+                        'line': prev_line_no + 1,
+                        'character': 0
+                    }
+                },
+                'newText': change.line + eol_chars
+            })
+        elif change.old:
+            # remove
+            lsp_line_no = change.old - 1
+            textEdits.append({
+                'range': {
+                    'start': {
+                        'line': lsp_line_no,
+                        'character': 0
+                    },
+                    'end': {
+                        # From LSP spec:
+                        # If you want to specify a range that contains a line
+                        # including the line ending character(s) then use an
+                        # end position denoting the start of the next line.
+                        'line': lsp_line_no + 1,
+                        'character': 0
+                    }
+                },
+                'newText': ''
+            })
+            prev_line_no = lsp_line_no
+
+    return textEdits

setup.cfg

@@ -33,6 +33,7 @@ all =
     pylint>=2.5.0
     rope>=0.10.5
     yapf
+    whatthepatch
 autopep8 = autopep8>=1.6.0,<1.7.0
 flake8 = flake8>=4.0.0,<4.1.0
 mccabe = mccabe>=0.6.0,<0.7.0
@@ -41,7 +42,9 @@ pydocstyle = pydocstyle>=2.0.0
 pyflakes = pyflakes>=2.4.0,<2.5.0
 pylint = pylint>=2.5.0
 rope = rope>0.10.5
-yapf = yapf
+yapf =
+    yapf
+    whatthepatch
 test =
     pylint>=2.5.0
     pytest

setup.py

@@ -1,3 +1,4 @@
+
 #!/usr/bin/env python3
 
 # Copyright 2017-2020 Palantir Technologies, Inc.
@@ -9,4 +10,4 @@
     setup(
         name="python-lsp-server",  # to allow GitHub dependency tracking work
         packages=find_packages(exclude=["contrib", "docs", "test", "test.*"]),  # https://github.com/pypa/setuptools/issues/2688
-    )
+    )

test/plugins/test_yapf_format.py

@@ -25,8 +25,7 @@ def test_format(workspace):
     doc = Document(DOC_URI, workspace, DOC)
     res = pylsp_format_document(doc)
 
-    assert len(res) == 1
-    assert res[0]['newText'] == "A = ['h', 'w', 'a']\n\nB = ['h', 'w']\n"
+    assert doc.apply_text_edits(res) == "A = ['h', 'w', 'a']\n\nB = ['h', 'w']\n"
 
 
 def test_range_format(workspace):
@@ -38,10 +37,8 @@ def test_range_format(workspace):
     }
     res = pylsp_format_range(doc, def_range)
 
-    assert len(res) == 1
-
     # Make sure B is still badly formatted
-    assert res[0]['newText'] == "A = ['h', 'w', 'a']\n\nB = ['h',\n\n\n'w']\n"
+    assert doc.apply_text_edits(res) == "A = ['h', 'w', 'a']\n\nB = ['h',\n\n\n'w']\n"
 
 
 def test_no_change(workspace):
@@ -56,12 +53,13 @@ def test_config_file(tmpdir, workspace):
     src = tmpdir.join('test.py')
     doc = Document(uris.from_fs_path(src.strpath), workspace, DOC)
 
-    # A was split on multiple lines because of column_limit from config file
-    assert pylsp_format_document(doc)[0]['newText'] == "A = [\n    'h', 'w',\n    'a'\n]\n\nB = ['h', 'w']\n"
+    res = pylsp_format_document(doc)
 
+    # A was split on multiple lines because of column_limit from config file
+    assert doc.apply_text_edits(res) == "A = [\n    'h', 'w',\n    'a'\n]\n\nB = ['h', 'w']\n"
 
 def test_cr_line_endings(workspace):
     doc = Document(DOC_URI, workspace, 'import os;import sys\r\rdict(a=1)')
     res = pylsp_format_document(doc)
 
-    assert res[0]['newText'] == 'import os\rimport sys\r\rdict(a=1)\r'
+    assert doc.apply_text_edits(res) == 'import os\rimport sys\r\rdict(a=1)\r'