Handle textDocument/foldingRange call (#665)

andfoy · ccordoba12 · commit 1ea46aee5491 · 2019-10-28T17:22:35.000-05:00
diff --git a/pyls/hookspecs.py b/pyls/hookspecs.py
@@ -67,6 +67,11 @@ def pyls_experimental_capabilities(config, workspace):
     pass
 
 
+@hookspec(firstresult=True)
+def pyls_folding_range(config, workspace, document):
+    pass
+
+
 @hookspec(firstresult=True)
 def pyls_format_document(config, workspace, document):
     pass
diff --git a/pyls/plugins/folding.py b/pyls/plugins/folding.py
@@ -0,0 +1,169 @@
+# pylint: disable=len-as-condition
+# Copyright 2019 Palantir Technologies, Inc.
+
+import re
+
+import parso
+import parso.python.tree as tree_nodes
+
+from pyls import hookimpl
+
+SKIP_NODES = (tree_nodes.Module, tree_nodes.IfStmt, tree_nodes.TryStmt)
+IDENTATION_REGEX = re.compile(r'(\s+).+')
+
+
+@hookimpl
+def pyls_folding_range(document):
+    program = document.source + '\n'
+    lines = program.splitlines()
+    tree = parso.parse(program)
+    ranges = __compute_folding_ranges(tree, lines)
+
+    results = []
+    for (start_line, end_line) in ranges:
+        start_line -= 1
+        end_line -= 1
+        # If start/end character is not defined, then it defaults to the
+        # corresponding line last character
+        results.append({
+            'startLine': start_line,
+            'endLine': end_line,
+        })
+    return results
+
+
+def __merge_folding_ranges(left, right):
+    for start in list(left.keys()):
+        right_start = right.pop(start, None)
+        if right_start is not None:
+            left[start] = max(right_start, start)
+    left.update(right)
+    return left
+
+
+def __empty_identation_stack(identation_stack, level_limits,
+                             current_line, folding_ranges):
+    while identation_stack != []:
+        upper_level = identation_stack.pop(0)
+        level_start = level_limits.pop(upper_level)
+        folding_ranges.append((level_start, current_line))
+    return folding_ranges
+
+
+def __match_identation_stack(identation_stack, level, level_limits,
+                             folding_ranges, current_line):
+    upper_level = identation_stack.pop(0)
+    while upper_level >= level:
+        level_start = level_limits.pop(upper_level)
+        folding_ranges.append((level_start, current_line))
+        upper_level = identation_stack.pop(0)
+    identation_stack.insert(0, upper_level)
+    return identation_stack, folding_ranges
+
+
+def __compute_folding_ranges_identation(text):
+    lines = text.splitlines()
+    folding_ranges = []
+    identation_stack = []
+    level_limits = {}
+    current_level = 0
+    current_line = 0
+    while lines[current_line] == '':
+        current_line += 1
+    for i, line in enumerate(lines):
+        if i < current_line:
+            continue
+        i += 1
+        identation_match = IDENTATION_REGEX.match(line)
+        if identation_match is not None:
+            whitespace = identation_match.group(1)
+            level = len(whitespace)
+            if level > current_level:
+                level_limits[current_level] = current_line
+                identation_stack.insert(0, current_level)
+                current_level = level
+            elif level < current_level:
+                identation_stack, folding_ranges = __match_identation_stack(
+                    identation_stack, level, level_limits, folding_ranges,
+                    current_line)
+                current_level = level
+        else:
+            folding_ranges = __empty_identation_stack(
+                identation_stack, level_limits, current_line, folding_ranges)
+            current_level = 0
+        if line.strip() != '':
+            current_line = i
+    folding_ranges = __empty_identation_stack(
+        identation_stack, level_limits, current_line, folding_ranges)
+    return dict(folding_ranges)
+
+
+def __check_if_node_is_valid(node):
+    valid = True
+    if isinstance(node, tree_nodes.PythonNode):
+        kind = node.type
+        valid = kind not in {'decorated', 'parameters'}
+        if kind == 'suite':
+            if isinstance(node.parent, tree_nodes.Function):
+                valid = False
+    return valid
+
+
+def __compute_start_end_lines(node, stack):
+    start_line, _ = node.start_pos
+    end_line, _ = node.end_pos
+
+    last_leaf = node.get_last_leaf()
+    last_newline = isinstance(last_leaf, tree_nodes.Newline)
+    last_operator = isinstance(last_leaf, tree_nodes.Operator)
+    node_is_operator = isinstance(node, tree_nodes.Operator)
+    last_operator = last_operator or not node_is_operator
+
+    end_line -= 1
+
+    modified = False
+    if isinstance(node.parent, tree_nodes.PythonNode):
+        kind = node.type
+        if kind in {'suite', 'atom', 'atom_expr', 'arglist'}:
+            if len(stack) > 0:
+                next_node = stack[0]
+                next_line, _ = next_node.start_pos
+                if next_line > end_line:
+                    end_line += 1
+                    modified = True
+    if not last_newline and not modified and not last_operator:
+        end_line += 1
+    return start_line, end_line
+
+
+def __compute_folding_ranges(tree, lines):
+    folding_ranges = {}
+    stack = [tree]
+
+    while len(stack) > 0:
+        node = stack.pop(0)
+        if isinstance(node, tree_nodes.Newline):
+            # Skip newline nodes
+            continue
+        elif isinstance(node, tree_nodes.PythonErrorNode):
+            # Fallback to identation-based (best-effort) folding
+            start_line, _ = node.start_pos
+            start_line -= 1
+            padding = [''] * start_line
+            text = '\n'.join(padding + lines[start_line:]) + '\n'
+            identation_ranges = __compute_folding_ranges_identation(text)
+            folding_ranges = __merge_folding_ranges(
+                folding_ranges, identation_ranges)
+            break
+        elif not isinstance(node, SKIP_NODES):
+            valid = __check_if_node_is_valid(node)
+            if valid:
+                start_line, end_line = __compute_start_end_lines(node, stack)
+                if end_line > start_line:
+                    current_end = folding_ranges.get(start_line, -1)
+                    folding_ranges[start_line] = max(current_end, end_line)
+        if hasattr(node, 'children'):
+            stack = node.children + stack
+
+    folding_ranges = sorted(folding_ranges.items())
+    return folding_ranges
diff --git a/pyls/python_ls.py b/pyls/python_ls.py
@@ -173,6 +173,7 @@ def capabilities(self):
             'hoverProvider': True,
             'referencesProvider': True,
             'renameProvider': True,
+            'foldingRangeProvider': True,
             'signatureHelpProvider': {
                 'triggerCharacters': ['(', ',', '=']
             },
@@ -282,6 +283,9 @@ def rename(self, doc_uri, position, new_name):
     def signature_help(self, doc_uri, position):
         return self._hook('pyls_signature_help', doc_uri, position=position)
 
+    def folding(self, doc_uri):
+        return self._hook('pyls_folding_range', doc_uri)
+
     def m_text_document__did_close(self, textDocument=None, **_kwargs):
         workspace = self._match_uri_to_workspace(textDocument['uri'])
         workspace.rm_document(textDocument['uri'])
@@ -333,6 +337,9 @@ def m_text_document__formatting(self, textDocument=None, _options=None, **_kwarg
     def m_text_document__rename(self, textDocument=None, position=None, newName=None, **_kwargs):
         return self.rename(textDocument['uri'], position, newName)
 
+    def m_text_document__folding_range(self, textDocument=None, **_kwargs):
+        return self.folding(textDocument['uri'])
+
     def m_text_document__range_formatting(self, textDocument=None, range=None, _options=None, **_kwargs):
         # Again, we'll ignore formatting options for now.
         return self.format_range(textDocument['uri'], range)
diff --git a/setup.py b/setup.py
@@ -78,6 +78,7 @@
         ],
         'pyls': [
             'autopep8 = pyls.plugins.autopep8_format',
+            'folding = pyls.plugins.folding',
             'flake8 = pyls.plugins.flake8_lint',
             'jedi_completion = pyls.plugins.jedi_completion',
             'jedi_definition = pyls.plugins.definition',
@@ -94,7 +95,7 @@
             'pylint = pyls.plugins.pylint_lint',
             'rope_completion = pyls.plugins.rope_completion',
             'rope_rename = pyls.plugins.rope_rename',
-            'yapf = pyls.plugins.yapf_format',
+            'yapf = pyls.plugins.yapf_format'
         ]
     },
 )
diff --git a/test/plugins/test_folding.py b/test/plugins/test_folding.py
@@ -0,0 +1,164 @@
+# Copyright 2019 Palantir Technologies, Inc.
+
+from textwrap import dedent
+
+from pyls import uris
+from pyls.workspace import Document
+from pyls.plugins.folding import pyls_folding_range
+
+
+DOC_URI = uris.from_fs_path(__file__)
+DOC = dedent("""
+def func(arg1, arg2, arg3,
+         arg4, arg5, default=func(
+             2, 3, 4
+         )):
+    return (2, 3,
+            4, 5)
+
+@decorator(
+    param1,
+    param2
+)
+def decorated_func(x, y, z):
+    if x:
+        return y
+    elif y:
+        return z
+    elif x + y > z:
+        return True
+    else:
+        return x
+
+class A():
+    def method(self, x1):
+        def inner():
+            return x1
+
+        if x2:
+            func(3, 4, 5, 6,
+                 7)
+        elif x3 < 2:
+            pass
+        else:
+            more_complex_func(2, 3, 4, 5, 6,
+                              8)
+        return inner
+
+a = 2
+operation = (a_large_variable_that_fills_all_space +
+             other_embarrasingly_long_variable - 2 * 3 / 5)
+
+(a, b, c,
+ d, e, f) = func(3, 4, 5, 6,
+                 7, 8, 9, 10)
+
+for i in range(0, 3):
+    i += 1
+    while x < i:
+        expr = (2, 4)
+        a = func(expr + i, arg2, arg3, arg4,
+                 arg5, var(2, 3, 4,
+                           5))
+    for j in range(0, i):
+        if i % 2 == 1:
+            pass
+
+compren = [x for x in range(0, 3)
+           if x == 2]
+
+with open('doc', 'r') as f:
+    try:
+        f / 0
+    except:
+        pass
+    finally:
+        raise SomeException()
+""")
+
+SYNTAX_ERR = dedent("""
+def func(arg1, arg2, arg3,
+         arg4, arg5, default=func(
+             2, 3, 4
+         )):
+    return (2, 3,
+            4, 5)
+
+class A(:
+    pass
+
+a = 2
+operation = (a_large_variable_that_fills_all_space +
+             other_embarrasingly_long_variable - 2 * 3 /
+
+(a, b, c,
+ d, e, f) = func(3, 4, 5, 6,
+                 7, 8, 9, 10
+a = 2
+for i in range(0, 3)
+    i += 1
+    while x < i:
+        expr = (2, 4)
+        a = func(expr + i, arg2, arg3, arg4,
+                 arg5, var(2, 3, 4,
+                           5))
+    for j in range(0, i):
+        if i % 2 == 1:
+            pass
+""")
+
+
+def test_folding():
+    doc = Document(DOC_URI, DOC)
+    ranges = pyls_folding_range(doc)
+    expected = [{'startLine': 1, 'endLine': 6},
+                {'startLine': 2, 'endLine': 3},
+                {'startLine': 5, 'endLine': 6},
+                {'startLine': 8, 'endLine': 11},
+                {'startLine': 12, 'endLine': 20},
+                {'startLine': 13, 'endLine': 14},
+                {'startLine': 15, 'endLine': 16},
+                {'startLine': 17, 'endLine': 18},
+                {'startLine': 19, 'endLine': 20},
+                {'startLine': 22, 'endLine': 35},
+                {'startLine': 23, 'endLine': 35},
+                {'startLine': 24, 'endLine': 25},
+                {'startLine': 27, 'endLine': 29},
+                {'startLine': 28, 'endLine': 29},
+                {'startLine': 30, 'endLine': 31},
+                {'startLine': 32, 'endLine': 34},
+                {'startLine': 33, 'endLine': 34},
+                {'startLine': 38, 'endLine': 39},
+                {'startLine': 41, 'endLine': 43},
+                {'startLine': 42, 'endLine': 43},
+                {'startLine': 45, 'endLine': 54},
+                {'startLine': 47, 'endLine': 51},
+                {'startLine': 49, 'endLine': 51},
+                {'startLine': 50, 'endLine': 51},
+                {'startLine': 52, 'endLine': 54},
+                {'startLine': 53, 'endLine': 54},
+                {'startLine': 56, 'endLine': 57},
+                {'startLine': 59, 'endLine': 65},
+                {'startLine': 60, 'endLine': 61},
+                {'startLine': 62, 'endLine': 63},
+                {'startLine': 64, 'endLine': 65}]
+    assert ranges == expected
+
+
+def test_folding_syntax_error():
+    doc = Document(DOC_URI, SYNTAX_ERR)
+    ranges = pyls_folding_range(doc)
+    expected = [{'startLine': 1, 'endLine': 6},
+                {'startLine': 2, 'endLine': 3},
+                {'startLine': 5, 'endLine': 6},
+                {'startLine': 8, 'endLine': 9},
+                {'startLine': 12, 'endLine': 13},
+                {'startLine': 15, 'endLine': 17},
+                {'startLine': 16, 'endLine': 17},
+                {'startLine': 19, 'endLine': 28},
+                {'startLine': 21, 'endLine': 25},
+                {'startLine': 23, 'endLine': 25},
+                {'startLine': 24, 'endLine': 25},
+                {'startLine': 26, 'endLine': 28},
+                {'startLine': 27, 'endLine': 28}]
+    assert ranges == expected
