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

masad-frost · web-flow · commit a9d503c5a766 · 2022-06-06T12:34:39.000-05:00
diff --git a/pylsp/_utils.py b/pylsp/_utils.py
@@ -14,6 +14,7 @@
 JEDI_VERSION = jedi.__version__
 
 # Eol chars accepted by the LSP protocol
+# the ordering affects performance
 EOL_CHARS = ['\r\n', '\r', '\n']
 EOL_REGEX = re.compile(f'({"|".join(EOL_CHARS)})')
 
diff --git a/pylsp/config/config.py b/pylsp/config/config.py
@@ -81,6 +81,7 @@ def __init__(self, root_uri, init_opts, process_id, capabilities):
             if plugin is not None:
                 log.info("Loaded pylsp plugin %s from %s", name, plugin)
 
+        # pylint: disable=no-member
         for plugin_conf in self._pm.hook.pylsp_settings(config=self):
             self._plugin_settings = _utils.merge_dicts(self._plugin_settings, plugin_conf)
 
diff --git a/pylsp/plugins/yapf_format.py b/pylsp/plugins/yapf_format.py
@@ -7,6 +7,8 @@
 from yapf.yapflib import file_resources, style
 from yapf.yapflib.yapf_api import FormatCode
 
+import whatthepatch
+
 from pylsp import hookimpl
 from pylsp._utils import get_eol_chars
 
@@ -36,75 +38,164 @@ def pylsp_format_range(document, range, options=None):  # pylint: disable=redefi
     return _format(document, lines=lines, options=options)
 
 
-def _format(document, lines=None, options=None):
-    # Yapf doesn't work with CRLF/CR line endings, so we replace them by '\n'
-    # and restore them below.
-    replace_eols = False
-    source = document.source
-    eol_chars = get_eol_chars(source)
-    if eol_chars in ['\r', '\r\n']:
-        replace_eols = True
-        source = source.replace(eol_chars, '\n')
-
+def get_style_config(document_path, options=None):
     # Get the default styles as a string
     # for a preset configuration, i.e. "pep8"
     style_config = file_resources.GetDefaultStyleForDir(
-        os.path.dirname(document.path)
+        os.path.dirname(document_path)
     )
-    if options is not None:
-        # We have options passed from LSP format request
-        # let's pass them to the formatter.
-        # First we want to get a dictionary of the preset style
-        # to pass instead of a string so that we can modify it
-        style_config = style.CreateStyleFromConfig(style_config)
-
-        use_tabs = style_config['USE_TABS']
-        indent_width = style_config['INDENT_WIDTH']
-
-        if options.get('tabSize') is not None:
-            indent_width = max(int(options.get('tabSize')), 1)
-
-        if options.get('insertSpaces') is not None:
-            # TODO is it guaranteed to be a boolean, or can it be a string?
-            use_tabs = not options.get('insertSpaces')
-
-            if use_tabs:
-                # Indent width doesn't make sense when using tabs
-                # the specifications state: "Size of a tab in spaces"
-                indent_width = 1
+    if options is None:
+        return style_config
+
+    # We have options passed from LSP format request
+    # let's pass them to the formatter.
+    # First we want to get a dictionary of the preset style
+    # to pass instead of a string so that we can modify it
+    style_config = style.CreateStyleFromConfig(style_config)
+
+    use_tabs = style_config['USE_TABS']
+    indent_width = style_config['INDENT_WIDTH']
+
+    if options.get('tabSize') is not None:
+        indent_width = max(int(options.get('tabSize')), 1)
+
+    if options.get('insertSpaces') is not None:
+        # TODO is it guaranteed to be a boolean, or can it be a string?
+        use_tabs = not options.get('insertSpaces')
+
+        if use_tabs:
+            # Indent width doesn't make sense when using tabs
+            # the specifications state: "Size of a tab in spaces"
+            indent_width = 1
+
+    style_config['USE_TABS'] = use_tabs
+    style_config['INDENT_WIDTH'] = indent_width
+    style_config['CONTINUATION_INDENT_WIDTH'] = indent_width
+
+    for style_option, value in options.items():
+        # Apply arbitrary options passed as formatter options
+        if style_option not in style_config:
+            # ignore if it's not a known yapf config
+            continue
+
+        style_config[style_option] = value
+
+    return style_config
+
+
+def diff_to_text_edits(diff, eol_chars):
+    # 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.
+    text_edits = []
+    # 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:
+            # old and new are the same line, no change
+            # diffs are 1-indexed
+            prev_line_no = change.old - 1
+        elif change.new:
+            # addition
+            text_edits.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
+            text_edits.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 text_edits
+
+
+def ensure_eof_new_line(document, eol_chars, text_edits):
+    # 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 document.source.endswith(eol_chars):
+        return
+
+    lines = document.lines
+    last_line_number = len(lines) - 1
+
+    if text_edits and text_edits[-1]['range']['start']['line'] >= last_line_number:
+        return
+
+    text_edits.append({
+        'range': {
+            'start': {
+                'line': last_line_number,
+                'character': 0
+            },
+            'end': {
+                'line': last_line_number + 1,
+                'character': 0
+            }
+        },
+        'newText': lines[-1] + eol_chars
+    })
 
-        style_config['USE_TABS'] = use_tabs
-        style_config['INDENT_WIDTH'] = indent_width
-        style_config['CONTINUATION_INDENT_WIDTH'] = indent_width
 
-        for style_option, value in options.items():
-            # Apply arbitrary options passed as formatter options
-            if style_option not in style_config:
-                # ignore if it's not a known yapf config
-                continue
+def _format(document, lines=None, options=None):
+    source = document.source
+    # Yapf doesn't work with CRLF/CR line endings, so we replace them by '\n'
+    # and restore them below when adding new lines
+    eol_chars = get_eol_chars(source)
+    if eol_chars in ['\r', '\r\n']:
+        source = source.replace(eol_chars, '\n')
+    else:
+        eol_chars = '\n'
 
-            style_config[style_option] = value
+    style_config = get_style_config(document_path=document.path, options=options)
 
-    new_source, changed = FormatCode(
+    diff_txt, changed = FormatCode(
         source,
         lines=lines,
         filename=document.filename,
+        print_diff=True,
         style_config=style_config
     )
 
     if not changed:
         return []
 
-    if replace_eols:
-        new_source = new_source.replace('\n', eol_chars)
+    patch_generator = whatthepatch.parse_patch(diff_txt)
+    diff = next(patch_generator)
+    patch_generator.close()
 
-    # 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
-    }]
+    text_edits = diff_to_text_edits(diff=diff, eol_chars=eol_chars)
+
+    ensure_eof_new_line(document=document, eol_chars=eol_chars, text_edits=text_edits)
+
+    return text_edits
diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
@@ -34,6 +34,7 @@ class _StreamHandlerWrapper(socketserver.StreamRequestHandler):
 
     def setup(self):
         super().setup()
+        # pylint: disable=no-member
         self.delegate = self.DELEGATE_CLASS(self.rfile, self.wfile)
 
     def handle(self):
@@ -47,6 +48,7 @@ def handle(self):
                 if isinstance(e, WindowsError) and e.winerror == 10054:
                     pass
 
+        # pylint: disable=no-member
         self.SHUTDOWN_CALL()
 
 
@@ -121,6 +123,7 @@ async def pylsp_ws(websocket):
             async for message in websocket:
                 try:
                     log.debug("consuming payload and feeding it to LSP handler")
+                    # pylint: disable=c-extension-no-member
                     request = json.loads(message)
                     loop = asyncio.get_running_loop()
                     await loop.run_in_executor(tpool, pylsp_handler.consume, request)
@@ -130,6 +133,7 @@ async def pylsp_ws(websocket):
         def send_message(message, websocket):
             """Handler to send responses of  processed requests to respective web socket clients"""
             try:
+                # pylint: disable=c-extension-no-member
                 payload = json.dumps(message, ensure_ascii=False)
                 asyncio.run(websocket.send(payload))
             except Exception as e:  # pylint: disable=broad-except
diff --git a/pylsp/text_edit.py b/pylsp/text_edit.py
@@ -0,0 +1,94 @@
+# Copyright 2017-2020 Palantir Technologies, Inc.
+# Copyright 2021- Python Language Server Contributors.
+
+def get_well_formatted_range(lsp_range):
+    start = lsp_range['start']
+    end = lsp_range['end']
+
+    if start['line'] > end['line'] or (start['line'] == end['line'] and start['character'] > end['character']):
+        return {'start': end, 'end': start}
+
+    return lsp_range
+
+
+def get_well_formatted_edit(text_edit):
+    lsp_range = get_well_formatted_range(text_edit['range'])
+    if lsp_range != text_edit['range']:
+        return {'newText': text_edit['newText'], 'range': lsp_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
+
+
+class OverLappingTextEditException(Exception):
+    """
+    Text edits are expected to be sorted
+    and compressed instead of overlapping.
+    This error is raised when two edits
+    are overlapping.
+    """
+
+
+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 OverLappingTextEditException('overlapping edit')
+
+        if 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)
diff --git a/pyproject.toml b/pyproject.toml
@@ -35,6 +35,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"]
@@ -44,7 +45,7 @@ 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>=1.0.2,<2.0.0"]
 websockets = ["websockets>=10.3"]
 test = [
     "pylint>=2.5.0",
diff --git a/test/plugins/test_yapf_format.py b/test/plugins/test_yapf_format.py
diff --git a/test/test_text_edit.py b/test/test_text_edit.py
