Start work on BlankLinesVisitor

Author: knutwannheden

rewrite/rewrite/python/__init__.py

@@ -66,6 +66,7 @@
 
     # Formatter
     'AutoFormat',
+    'BlankLinesVisitor',
     'NormalizeFormatVisitor',
     'SpacesVisitor',
 ]

rewrite/rewrite/python/format/__init__.py

@@ -1,10 +1,12 @@
 __path__ = __import__('pkgutil').extend_path(__path__, __name__)
 
 from .auto_format import *
+from .blank_lines import *
 from .normalize_format import *
 
 __all__ = [
     'AutoFormat',
+    'BlankLinesVisitor',
     'NormalizeFormatVisitor',
     'SpacesVisitor',
 ]

rewrite/rewrite/python/format/auto_format.py

@@ -1,6 +1,7 @@
 from typing import Optional, cast, TypeVar
 
 from .normalize_format import NormalizeFormatVisitor
+from .. import BlankLinesVisitor, BlankLinesStyle
 from ..style import SpacesStyle, IntelliJ
 from ..visitor import PythonVisitor
 from ... import Recipe, Tree, Cursor
@@ -22,6 +23,7 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) ->
         cu = tree if isinstance(tree, JavaSourceFile) else self._cursor.first_enclosing_or_throw(JavaSourceFile)
 
         tree = NormalizeFormatVisitor(self._stop_after).visit(tree, p, self._cursor.fork())
+        tree = BlankLinesVisitor(cu.get_style(BlankLinesStyle) or IntelliJ.blank_lines(), self._stop_after).visit(tree, p, self._cursor.fork())
         tree = SpacesVisitor(cu.get_style(SpacesStyle) or IntelliJ.spaces(), self._stop_after).visit(tree, p, self._cursor.fork())
         return tree
 

rewrite/rewrite/python/format/blank_lines.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import Optional, TypeVar
+
+from rewrite import Tree, P, Cursor
+from rewrite.java import J
+from rewrite.python import PythonVisitor, BlankLinesStyle
+from rewrite.visitor import T
+
+J2 = TypeVar('J2', bound=J)
+
+
+class BlankLinesVisitor(PythonVisitor):
+    def __init__(self, style: BlankLinesStyle, stop_after: Tree = None):
+        self._style = style
+        self._stop_after = stop_after
+        self._stop = False
+
+    def post_visit(self, tree: T, p: P) -> Optional[T]:
+        if self._stop_after and tree == self._stop_after:
+            self._stop = True
+        return tree
+
+    def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> Optional[T]:
+        return tree if self._stop else super().visit(tree, p, parent)