diff --git a/pyls/plugins/folding.py b/pyls/plugins/folding.py index 48cfdfd5..7b44046c 100644 --- a/pyls/plugins/folding.py +++ b/pyls/plugins/folding.py @@ -102,16 +102,51 @@ def __check_if_node_is_valid(node): valid = True if isinstance(node, tree_nodes.PythonNode): kind = node.type - valid = kind not in {'decorated', 'parameters'} + 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_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', 'except'}: + body = stack[2] + children = [body] + if hasattr(body, 'children'): + children = body.children + stack = stack[:2] + children + stack[3:] + node = body + end_line, _ = body.end_pos + elif node.value in {'for'}: + body = stack[4] + children = [body] + if hasattr(body, 'children'): + children = body.children + stack = stack[:4] + children + stack[5:] + node = body + end_line, _ = body.end_pos + elif node.value in {'else'}: + body = stack[1] + children = [body] + if hasattr(body, 'children'): + children = body.children + stack = stack[:1] + children + stack[2:] + node = body + end_line, _ = body.end_pos + 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) @@ -121,8 +156,7 @@ def __compute_start_end_lines(node, stack): end_line -= 1 - modified = False - if isinstance(node.parent, tree_nodes.PythonNode): + 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: @@ -133,7 +167,7 @@ def __compute_start_end_lines(node, stack): modified = True if not last_newline and not modified and not last_operator: end_line += 1 - return start_line, end_line + return start_line, end_line, stack def __compute_folding_ranges(tree, lines): @@ -158,7 +192,8 @@ def __compute_folding_ranges(tree, lines): 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) + 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)