Add blank linkes after group of top-level imports

Author: knutwannheden

rewrite/rewrite/python/_parser_visitor.py

@@ -584,7 +584,7 @@ def visit_Import(self, node):
             False,
             JContainer(
                 Space.EMPTY,
-                [self.__pad_list_element(self.__convert(n), i == len(node.names) - 1) for i, n in
+                [self.__pad_list_element(self.__convert(n), i == len(node.names) - 1, pad_last=False) for i, n in
                  enumerate(node.names)],
                 Markers.EMPTY
             )

rewrite/rewrite/python/format/blank_lines.py

@@ -3,8 +3,8 @@
 from typing import Optional, TypeVar, cast
 
 from rewrite import Tree, P, Cursor
-from rewrite.java import J, Space, Statement, JRightPadded, Block, ClassDeclaration, MethodDeclaration
-from rewrite.python import PythonVisitor, BlankLinesStyle, CompilationUnit
+from rewrite.java import J, Space, Statement, JRightPadded, Block, ClassDeclaration, MethodDeclaration, Import
+from rewrite.python import PythonVisitor, BlankLinesStyle, CompilationUnit, MultiImport
 from rewrite.visitor import T
 
 J2 = TypeVar('J2', bound=J)
@@ -26,10 +26,21 @@ def visit_statement(self, statement: Statement, p: P) -> J:
 
         parent_cursor = self.cursor.parent_tree_cursor()
         top_level = isinstance(parent_cursor.value, CompilationUnit)
-        if top_level and statement != cast(CompilationUnit, parent_cursor.value).statements[0]:
-            statement = minimum_lines_for_tree(statement, self._style.minimum.around_top_level_classes_functions)
-        elif top_level:
-            statement = statement.with_prefix(statement.prefix.with_whitespace(''))
+        if top_level and isinstance(statement, (Import, MultiImport)):
+            parent_cursor.put_message('previous_import', True)
+            prev_import = False
+        else:
+            prev_import = top_level and parent_cursor.get_message('previous_import', False)
+            if prev_import:
+                parent_cursor.put_message('previous_import', False)
+
+        if top_level:
+            if statement == cast(CompilationUnit, parent_cursor.value).statements[0]:
+                statement = statement.with_prefix(statement.prefix.with_whitespace(''))
+            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)
         else:
             in_block = isinstance(parent_cursor.value, Block)
             in_class = in_block and isinstance(parent_cursor.parent_tree_cursor().value, ClassDeclaration)

rewrite/rewrite/visitor.py

@@ -25,6 +25,11 @@ class Cursor:
     def get_message(self, key: str, default_value: O) -> O:
         return default_value if self.messages is None else cast(O, self.messages.get(key))
 
+    def put_message(self, key: str, value: object) -> None:
+        if self.messages is None:
+            object.__setattr__(self, 'messages', {})
+        self.messages[key] = value
+
     def parent_tree_cursor(self) -> Optional[Cursor]:
         c = self.parent
         while c is not None:

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

@@ -83,3 +83,34 @@ class Nested:
             from_visitor(BlankLinesVisitor(IntelliJ.blank_lines()))
         )
     )
+
+
+def test_blank_lines_after_top_level_imports():
+    style = IntelliJ.blank_lines()
+    style = style.with_minimum(
+        style.minimum.with_after_top_level_imports(3)
+    )
+    rewrite_run(
+        # language=python
+        python(
+            """\
+            import os
+            import sys
+            class Foo:
+                pass
+            """,
+            """\
+            import os
+            import sys
+
+
+
+            class Foo:
+                pass
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipes(
+            from_visitor(BlankLinesVisitor(style))
+        )
+    )