210 lines
6.9 KiB
Python
210 lines
6.9 KiB
Python
# Copyright 2017-2020 Palantir Technologies, Inc.
|
|
# Copyright 2021- Python Language Server Contributors.
|
|
|
|
import re
|
|
|
|
import parso
|
|
import parso.python.tree as tree_nodes
|
|
|
|
from pylsp import hookimpl
|
|
|
|
SKIP_NODES = (tree_nodes.Module, tree_nodes.IfStmt, tree_nodes.TryStmt)
|
|
IDENTATION_REGEX = re.compile(r"(\s+).+")
|
|
|
|
|
|
@hookimpl
|
|
def pylsp_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",
|
|
"dictorsetmaker",
|
|
"testlist_comp",
|
|
}
|
|
if kind == "suite":
|
|
if isinstance(node.parent, tree_nodes.Function):
|
|
valid = False
|
|
return valid
|
|
|
|
|
|
def __handle_skip(stack, skip):
|
|
body = stack[skip]
|
|
children = [body]
|
|
if hasattr(body, "children"):
|
|
children = body.children
|
|
stack = stack[:skip] + children + stack[skip + 1 :]
|
|
node = body
|
|
end_line, _ = body.end_pos
|
|
return node, end_line
|
|
|
|
|
|
def __handle_flow_nodes(node, end_line, stack):
|
|
from_keyword = False
|
|
if isinstance(node, tree_nodes.Keyword):
|
|
from_keyword = True
|
|
if node.value in {"if", "elif", "with", "while"}:
|
|
node, end_line = __handle_skip(stack, 2)
|
|
elif node.value in {"except"}:
|
|
first_node = stack[0]
|
|
if isinstance(first_node, tree_nodes.Operator):
|
|
node, end_line = __handle_skip(stack, 1)
|
|
else:
|
|
node, end_line = __handle_skip(stack, 2)
|
|
elif node.value in {"for"}:
|
|
node, end_line = __handle_skip(stack, 4)
|
|
elif node.value in {"else"}:
|
|
node, end_line = __handle_skip(stack, 1)
|
|
return end_line, from_keyword, node, stack
|
|
|
|
|
|
def __compute_start_end_lines(node, stack):
|
|
start_line, _ = node.start_pos
|
|
end_line, _ = node.end_pos
|
|
modified = False
|
|
end_line, from_keyword, node, stack = __handle_flow_nodes(node, end_line, stack)
|
|
|
|
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
|
|
|
|
if isinstance(node.parent, tree_nodes.PythonNode) and not from_keyword:
|
|
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, stack
|
|
|
|
|
|
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
|
|
if isinstance(node, tree_nodes.PythonErrorNode):
|
|
# Fallback to indentation-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
|
|
if not isinstance(node, SKIP_NODES):
|
|
valid = __check_if_node_is_valid(node)
|
|
if valid:
|
|
start_line, end_line, stack = __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
|