Implement support for blank_lines.keep_maximum styles

Author: knutwannheden

rewrite/rewrite/python/_parser_visitor.py

@@ -734,10 +734,7 @@ def visit_GeneratorExp(self, node):
         )
 
     def visit_Expr(self, node):
-        return py.ExpressionStatement(
-            random_id(),
-            self.__convert(node.value)
-        )
+        return self.__convert(node.value)
 
     def visit_Yield(self, node):
         return py.StatementExpression(
@@ -1941,6 +1938,15 @@ def __convert_type_mapper(self, node) -> Optional[TypeTree]:
     def __convert(self, node) -> Optional[J]:
         return self.__convert_internal(node, self.__convert)
 
+    def __convert_statement(self, node) -> Optional[J]:
+        converted = self.__convert_internal(node, self.__convert_statement)
+        if is_of_type(converted, Statement):
+            return converted
+        return py.ExpressionStatement(
+            random_id(),
+            converted
+        )
+
     def __convert_internal(self, node, recursion, mapping = None) -> Optional[J]:
         if not node or not isinstance(node, ast.expr) or isinstance(node, ast.GeneratorExp):
             return self.visit(cast(ast.AST, node)) if node else None
@@ -2065,7 +2071,7 @@ def __convert_block(self, statements: Sequence, prefix: str = ':') -> j.Block:
         )
 
     def __pad_statement(self, stmt: ast.stmt) -> JRightPadded[Statement]:
-        statement = self.__convert(stmt)
+        statement = self.__convert_statement(stmt)
         # use whitespace until end of line as padding; what follows will be the prefix of next element
         save_cursor = self._cursor
         padding = self.__whitespace('\n')

rewrite/rewrite/python/format/blank_lines.py

@@ -41,7 +41,7 @@ def visit_statement(self, statement: Statement, p: P) -> J:
             else:
                 min_lines = max(self._style.minimum.around_top_level_classes_functions if isinstance(statement, (ClassDeclaration, MethodDeclaration)) else 0,
                                 self._style.minimum.after_top_level_imports if prev_import else 0)
-                statement = minimum_lines_for_tree(statement, min_lines)
+                statement = _adjusted_lines_for_tree(statement, min_lines, self._style.keep_maximum.in_declarations)
         else:
             in_block = isinstance(parent_cursor.value, Block)
             in_class = in_block and isinstance(parent_cursor.parent_tree_cursor().value, ClassDeclaration)
@@ -54,9 +54,16 @@ def visit_statement(self, statement: Statement, p: P) -> J:
                     min_lines = max(min_lines, self._style.minimum.around_class)
                 elif is_first and isinstance(statement, MethodDeclaration):
                     min_lines = max(min_lines, self._style.minimum.before_first_method)
+
+            # This seems to correspond to how IntelliJ interprets this configuration
+            max_lines = self._style.keep_maximum.in_declarations if \
+                isinstance(statement, (ClassDeclaration, MethodDeclaration)) else \
+                self._style.keep_maximum.in_code
+
             if prev_import:
                 min_lines = max(min_lines, self._style.minimum.after_local_imports)
-            statement = minimum_lines_for_tree(statement, min_lines)
+
+            statement = _adjusted_lines_for_tree(statement, min_lines, max_lines)
         return statement
 
     def post_visit(self, tree: T, p: P) -> Optional[T]:
@@ -68,39 +75,35 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) ->
         return tree if self._stop else super().visit(tree, p, parent)
 
 
-def minimum_lines_for_right_padded(tree: JRightPadded[J2], min_lines) -> JRightPadded[J2]:
-    return tree.with_element(minimum_lines_for_tree(tree.element, min_lines))
+def _adjusted_lines_for_right_padded(tree: JRightPadded[J2], min_lines: int, max_lines: int) -> JRightPadded[J2]:
+    return tree.with_element(_adjusted_lines_for_tree(tree.element, min_lines, max_lines))
 
 
-def minimum_lines_for_tree(tree: J, min_lines) -> J:
-    if min_lines == 0:
-        return tree
-    return tree.with_prefix(minimum_lines_for_space(tree.prefix, min_lines))
+def _adjusted_lines_for_tree(tree: J, min_lines: int, max_lines: int) -> J:
+    return tree.with_prefix(_adjusted_lines_for_space(tree.prefix, min_lines, max_lines))
 
 
-def minimum_lines_for_space(prefix: Space, min_lines) -> Space:
-    if min_lines == 0:
-        return prefix
+def _adjusted_lines_for_space(prefix: Space, min_lines: int, max_lines: int) -> Space:
     if not prefix.comments or \
             '\n' in prefix.whitespace or \
             (prefix.comments[0].multiline and '\n' in prefix.comments[0].text):
-        return prefix.with_whitespace(minimum_lines_for_string(prefix.whitespace, min_lines))
+        return prefix.with_whitespace(_adjusted_lines_for_string(prefix.whitespace, min_lines, max_lines))
 
     # the first comment is a trailing comment on the previous line
-    c0 = prefix.comments[0].with_suffix(minimum_lines_for_string(prefix.comments[0].suffix, min_lines))
+    c0 = prefix.comments[0].with_suffix(_adjusted_lines_for_string(prefix.comments[0].suffix, min_lines, max_lines))
     return prefix if c0 is prefix.comments[0] else prefix.with_comments([c0] + prefix.comments[1:])
 
 
-def minimum_lines_for_string(whitespace, min_lines):
-    if min_lines == 0:
+def _adjusted_lines_for_string(whitespace, min_lines: int, max_lines: int):
+    existing_blank_lines = max(_count_line_breaks(whitespace) - 1, 0)
+    max_lines = max(min_lines, max_lines)
+    if min_lines <= existing_blank_lines <= max_lines:
         return whitespace
-
-    min_whitespace = whitespace
-    for _ in range(min_lines - get_new_line_count(whitespace) + 1):
-        min_whitespace = '\n' + min_whitespace
-
-    return min_whitespace
+    elif existing_blank_lines < min_lines:
+        return '\n' * (min_lines - existing_blank_lines) + whitespace
+    else:
+        return '\n' * max_lines + whitespace[whitespace.rfind('\n'):]
 
 
-def get_new_line_count(whitespace):
+def _count_line_breaks(whitespace):
     return whitespace.count('\n')

rewrite/tests/python/all/format/blank_lines_test.py

@@ -187,3 +187,32 @@ def __init__(self):
             from_visitor(BlankLinesVisitor(style))
         )
     )
+
+
+def test_max_lines_in_code():
+    style = IntelliJ.blank_lines()
+    style = style.with_keep_maximum(
+        style.keep_maximum.with_in_code(1)
+    )
+    rewrite_run(
+        # language=python
+        python(
+            """\
+            def foo():
+                print('foo')
+
+
+                print('bar')
+            """,
+            """\
+            def foo():
+                print('foo')
+
+                print('bar')
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipes(
+            from_visitor(BlankLinesVisitor(style))
+        )
+    )