Add first BlankLinesVisitor test

Author: knutwannheden

rewrite/rewrite/python/format/auto_format.py

@@ -1,8 +1,8 @@
 from typing import Optional, cast, TypeVar
 
+from .blank_lines import BlankLinesVisitor
 from .normalize_format import NormalizeFormatVisitor
-from .. import BlankLinesVisitor, BlankLinesStyle
-from ..style import SpacesStyle, IntelliJ
+from ..style import BlankLinesStyle, SpacesStyle, IntelliJ
 from ..visitor import PythonVisitor
 from ... import Recipe, Tree, Cursor
 from ...java import JavaSourceFile, MethodDeclaration, J, Space

rewrite/rewrite/python/format/blank_lines.py

@@ -3,8 +3,8 @@
 from typing import Optional, TypeVar
 
 from rewrite import Tree, P, Cursor
-from rewrite.java import J
-from rewrite.python import PythonVisitor, BlankLinesStyle
+from rewrite.java import J, Space, Statement
+from rewrite.python import PythonVisitor, BlankLinesStyle, CompilationUnit
 from rewrite.visitor import T
 
 J2 = TypeVar('J2', bound=J)
@@ -16,6 +16,15 @@ def __init__(self, style: BlankLinesStyle, stop_after: Tree = None):
         self._stop_after = stop_after
         self._stop = False
 
+    def visit_compilation_unit(self, compilation_unit: CompilationUnit, p: P) -> J:
+        if not compilation_unit.prefix.comments:
+            compilation_unit = compilation_unit.with_prefix(Space.EMPTY)
+        return super().visit_compilation_unit(compilation_unit, p)
+
+    def visit_statement(self, statement: Statement, p: P) -> J:
+        statement = statement.with_prefix(Space.EMPTY)
+        return super().visit_statement(statement, p)
+
     def post_visit(self, tree: T, p: P) -> Optional[T]:
         if self._stop_after and tree == self._stop_after:
             self._stop = True

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

@@ -0,0 +1,21 @@
+from rewrite.java import Space
+from rewrite.python import IntelliJ
+from rewrite.python.format import BlankLinesVisitor
+from rewrite.test import rewrite_run, python, RecipeSpec, from_visitor
+
+
+def test_remove_leading_module_blank_lines():
+    rewrite_run(
+        # language=python
+        python(
+            """
+
+
+            print('foo')
+            """,
+            """print('foo')
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipe(from_visitor(BlankLinesVisitor(IntelliJ.blank_lines())))
+    )