Add delegate prefix to ExpressionStatement

Author: knutwannheden

rewrite/rewrite/python/format/blank_lines.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional, TypeVar
+from typing import Optional, TypeVar, cast
 
 from rewrite import Tree, P, Cursor
 from rewrite.java import J, Space, Statement
@@ -22,8 +22,15 @@ def visit_compilation_unit(self, compilation_unit: CompilationUnit, p: P) -> J:
         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)
+        statement = super().visit_statement(statement, p)
+
+        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 = statement.with_prefix(statement.prefix.with_whitespace('\n\n'))
+        elif top_level:
+            statement = statement.with_prefix(statement.prefix.with_whitespace(''))
+        return statement
 
     def post_visit(self, tree: T, p: P) -> Optional[T]:
         if self._stop_after and tree == self._stop_after:

rewrite/rewrite/python/tree.py

@@ -719,6 +719,13 @@ def id(self) -> UUID:
     def with_id(self, id: UUID) -> ExpressionStatement:
         return self if id is self._id else replace(self, _id=id)
 
+    @property
+    def prefix(self) -> Space:
+        return self._expression.prefix
+
+    def with_prefix(self, prefix: Space) -> ExpressionStatement:
+        return self.with_expression(self._expression.with_prefix(prefix))
+
     _expression: Expression
 
     @property
@@ -790,7 +797,7 @@ def prefix(self) -> Space:
         return self._statement.prefix
 
     def with_prefix(self, prefix: Space) -> StatementExpression:
-        return self._statement.with_prefix(prefix)
+        return self.with_statement(self._statement.with_prefix(prefix))
 
     _statement: Statement
 

rewrite/rewrite/python/visitor.py

@@ -98,6 +98,14 @@ def visit_compilation_unit(self, compilation_unit: CompilationUnit, p: P) -> J:
         return compilation_unit
 
     def visit_expression_statement(self, expression_statement: ExpressionStatement, p: P) -> J:
+        temp_statement = cast(Statement, self.visit_statement(expression_statement, p))
+        if not isinstance(temp_statement, ExpressionStatement):
+            return temp_statement
+        expression_statement = cast(ExpressionStatement, temp_statement)
+        temp_expression = cast(Expression, self.visit_expression(expression_statement, p))
+        if not isinstance(temp_expression, ExpressionStatement):
+            return temp_expression
+        expression_statement = cast(ExpressionStatement, temp_expression)
         expression_statement = expression_statement.with_expression(self.visit_and_cast(expression_statement.expression, Expression, p))
         return expression_statement
 
@@ -112,6 +120,14 @@ def visit_expression_type_tree(self, expression_type_tree: ExpressionTypeTree, p
         return expression_type_tree
 
     def visit_statement_expression(self, statement_expression: StatementExpression, p: P) -> J:
+        temp_statement = cast(Statement, self.visit_statement(statement_expression, p))
+        if not isinstance(temp_statement, StatementExpression):
+            return temp_statement
+        statement_expression = cast(StatementExpression, temp_statement)
+        temp_expression = cast(Expression, self.visit_expression(statement_expression, p))
+        if not isinstance(temp_expression, StatementExpression):
+            return temp_expression
+        statement_expression = cast(StatementExpression, temp_expression)
         statement_expression = statement_expression.with_statement(self.visit_and_cast(statement_expression.statement, Statement, p))
         return statement_expression
 

rewrite/rewrite/visitor.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Protocol, TypeVar, Optional, Dict, List, Any, cast, Type, ClassVar, TYPE_CHECKING
+from typing import Protocol, TypeVar, Optional, Dict, List, Any, cast, Type, ClassVar
 
 from .execution import RecipeRunException
 from .markers import Marker, Markers
@@ -25,6 +25,14 @@ 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 parent_tree_cursor(self) -> Optional[Cursor]:
+        c = self.parent
+        while c is not None:
+            if isinstance(c.value, Tree):
+                return c
+            c = c.parent
+        return None
+
     def first_enclosing_or_throw(self, type: Type[P]) -> P:
         result = self.first_enclosing(type)
         if result is None: