Allow source spec to include after recipe

Author: knutwannheden

rewrite/rewrite/python/__init__.py

@@ -27,6 +27,7 @@
     'Await',
     'ChainedAssignment',
     'CollectionLiteral',
+    'CompilationUnit',
     'ComprehensionExpression',
     'Del',
     'DictLiteral',

rewrite/rewrite/test.py

@@ -5,15 +5,18 @@
 from dataclasses import dataclass, field
 from io import StringIO
 from pathlib import Path
-from typing import Optional, Callable, Iterable, List
+from typing import Optional, Callable, Iterable, List, TypeVar
 from uuid import UUID
 
 from rewrite import InMemoryExecutionContext, ParserInput, ParserBuilder, random_id, ParseError, ParseExceptionResult, \
-    ExecutionContext, Recipe, TreeVisitor
+    ExecutionContext, Recipe, TreeVisitor, SourceFile
 from rewrite.execution import InMemoryLargeSourceSet
+from rewrite.python import CompilationUnit
 from rewrite.python.parser import PythonParserBuilder
 
 
+S = TypeVar('S', bound=SourceFile)
+
 @dataclass(frozen=True, eq=False)
 class SourceSpec:
     _id: UUID
@@ -46,6 +49,12 @@ def after(self) -> Optional[Callable[[str], str]]:
     def source_path(self) -> Optional[Path]:
         return self._source_path
 
+    _after_recipe: Callable[[S], None] = lambda _: None
+
+    @property
+    def after_recipe(self) -> Callable[[S], None]:
+        return self._after_recipe
+
 
 @dataclass(frozen=True)
 class CompositeRecipe(Recipe):
@@ -118,6 +127,7 @@ def rewrite_run(*source_specs: Iterable[SourceSpec], spec: RecipeSpec = None):
             for res in result:
                 if res._before and res._after:
                     source_spec = spec_by_source_file[res._before]
+                    source_spec.after_recipe(res._after)
                     after_printed = res._after.print_all()
                     if source_spec.after is not None:
                         after = source_spec.after(after_printed)
@@ -131,13 +141,14 @@ def rewrite_run(*source_specs: Iterable[SourceSpec], spec: RecipeSpec = None):
         remoting_context.close()
 
 
-def python(before: str, after: str = None) -> list[SourceSpec]:
+def python(before: str, after: str = None, after_recipe: Callable[[CompilationUnit], None] = lambda s: None) -> list[SourceSpec]:
     return [SourceSpec(
         random_id(),
         PythonParserBuilder(),
         textwrap.dedent(before),
         None if after is None else lambda _: textwrap.dedent(after),
-        None
+        None,
+        after_recipe
     )]
 
 

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

@@ -1,4 +1,5 @@
-from rewrite.python import NormalizeFormatVisitor, PythonVisitor
+from rewrite.java import Space
+from rewrite.python import NormalizeFormatVisitor, PythonVisitor, CompilationUnit
 from rewrite.test import rewrite_run, python, RecipeSpec, from_visitor
 
 
@@ -8,6 +9,9 @@ def visit_method_declaration(self, method, p):
 
 
 def test_remove_decorator():
+    def assert_prefix(cu: CompilationUnit):
+        assert cu.statements[1].prefix == Space([], '\n')
+
     rewrite_run(
         # language=python
         python(
@@ -24,7 +28,8 @@ def f(n):
 
             def f(n):
                 return n
-            """
+            """,
+            assert_prefix
         ),
         spec=RecipeSpec()
         .with_recipes(