Parse YAPF diffs into TextEdits (instead of sending the full doc)

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

pylsp/workspace.py

@@ -129,6 +129,84 @@ def _create_document(self, doc_uri, source=None, version=None):
         )
 
 
+def get_well_formatted_range(range):
+    start = range['start']
+    end = range['end']
+
+    if start['line'] > end['line'] or (start['line'] == end['line'] and start['character'] > end['character']):
+        return { 'start': end, 'end': start }
+
+    return range
+
+def get_well_formatted_edit(text_edit):
+    range = get_well_formatted_range(text_edit['range'])
+    if range != text_edit['range']:
+        return { 'newText': text_edit['newText'], 'range': range }
+
+    return text_edit
+
+def compare_text_edits(a, b):
+    diff = a['range']['start']['line'] - b['range']['start']['line']
+    if diff == 0:
+        return a['range']['start']['character'] - b['range']['start']['character']
+
+    return diff
+
+def merge_sort_text_edits(text_edits):
+    if len(text_edits) <= 1:
+        return text_edits
+
+    p = len(text_edits) // 2
+    left = text_edits[:p]
+    right = text_edits[p:]
+
+    merge_sort_text_edits(left)
+    merge_sort_text_edits(right)
+
+    left_idx = 0
+    right_idx = 0
+    i = 0
+    while left_idx < len(left) and right_idx < len(right):
+        ret = compare_text_edits(left[left_idx], right[right_idx])
+        if ret <= 0:
+            #  smaller_equal -> take left to preserve order
+            text_edits[i] = left[left_idx]
+            i += 1
+            left_idx += 1
+        else:
+            # greater -> take right
+            text_edits[i] = right[right_idx]
+            i+=1
+            right_idx +=1
+    while left_idx < len(left):
+        text_edits[i] = left[left_idx]
+        i += 1
+        left_idx += 1
+    while right_idx < len(right):
+        text_edits[i] = right[right_idx]
+        i += 1
+        right_idx += 1
+    return text_edits
+
+def apply_text_edits(doc, text_edits):
+    text = doc.source
+    sorted_edits = merge_sort_text_edits(list(map(get_well_formatted_edit,text_edits)))
+    last_modified_offset = 0
+    spans = []
+    for e in sorted_edits:
+        start_offset = doc.offset_at_position(e['range']['start'])
+        if start_offset < last_modified_offset:
+            raise Exception('overlapping edit')
+        elif start_offset > last_modified_offset:
+            spans.append(text[last_modified_offset:start_offset])
+
+        if len(e['newText']):
+            spans.append(e['newText'])
+        last_modified_offset = doc.offset_at_position(e['range']['end'])
+
+    spans.append(text[last_modified_offset:])
+    return ''.join(spans)
+
 class Document:
 
     def __init__(self, uri, workspace, source=None, version=None, local=True, extra_sys_path=None,

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

test/plugins/test_yapf_format.py

@@ -3,7 +3,7 @@
 
 from pylsp import uris
 from pylsp.plugins.yapf_format import pylsp_format_document, pylsp_format_range
-from pylsp.workspace import Document
+from pylsp.workspace import Document, apply_text_edits
 
 DOC_URI = uris.from_fs_path(__file__)
 DOC = """A = [
@@ -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 apply_text_edits(doc, 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 apply_text_edits(doc, 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 apply_text_edits(doc, 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 apply_text_edits(doc, res) == 'import os\rimport sys\r\rdict(a=1)'