Python: Handle PrepareRecipe, Visit, and Generate (#6551)

Author: knutwannheden

.gitignore

@@ -31,3 +31,4 @@ dist/
 .pytest_cache/
 .mypy_cache/
 .ruff_cache/
+.venv/

rewrite-python/rewrite/src/rewrite/rpc/server.py

@@ -408,6 +408,10 @@ def handle_reset(params: dict) -> bool:
     local_objects.clear()
     remote_objects.clear()
     remote_refs.clear()
+    _prepared_recipes.clear()
+    _execution_contexts.clear()
+    _recipe_accumulators.clear()
+    _recipe_phases.clear()
     logger.info("Reset: cleared all cached state")
     return True
 
@@ -521,6 +525,281 @@ def _serialize_value(value) -> Any:
     return str(value)
 
 
+# Prepared recipes storage - maps recipe IDs to recipe instances
+_prepared_recipes: Dict[str, Any] = {}
+# Execution contexts storage - maps context IDs to ExecutionContext instances
+_execution_contexts: Dict[str, Any] = {}
+# Accumulator storage for ScanningRecipes - maps recipe IDs to accumulators
+_recipe_accumulators: Dict[str, Any] = {}
+# Phase tracking for recipes - maps recipe IDs to 'scan' or 'edit'
+_recipe_phases: Dict[str, str] = {}
+
+
+def handle_prepare_recipe(params: dict) -> dict:
+    """Handle a PrepareRecipe RPC request.
+
+    Prepares a recipe for execution by:
+    1. Looking up the recipe in the marketplace
+    2. Instantiating it with the provided options
+    3. Storing it with a unique ID
+    4. Returning the descriptor and visitor info
+
+    Args:
+        params: dict with 'id' (recipe name) and optional 'options'
+
+    Returns:
+        dict with 'id', 'descriptor', 'editVisitor', and precondition info
+    """
+    recipe_name = params.get('id')
+    if recipe_name is None:
+        raise ValueError("Recipe 'id' is required")
+    options = params.get('options', {})
+
+    logger.info(f"PrepareRecipe: id={recipe_name}, options={options}")
+
+    marketplace = _get_marketplace()
+
+    # Look up the recipe - returns (RecipeDescriptor, Type[Recipe]) tuple
+    recipe_info = marketplace.find_recipe(recipe_name)
+    if recipe_info is None:
+        raise ValueError(f"Recipe not found: {recipe_name}")
+
+    _descriptor, recipe_class = recipe_info
+    if recipe_class is None:
+        raise ValueError(f"Recipe class not found for: {recipe_name}")
+
+    # Instantiate the recipe with options
+    if options:
+        recipe = recipe_class(**options)
+    else:
+        recipe = recipe_class()
+
+    # Generate a unique ID for this prepared recipe
+    prepared_id = generate_id()
+    _prepared_recipes[prepared_id] = recipe
+
+    # Build the response
+    descriptor = recipe.descriptor()
+
+    # Determine if this is a scanning recipe
+    from rewrite.recipe import ScanningRecipe
+    is_scanning = isinstance(recipe, ScanningRecipe)
+
+    response = {
+        'id': prepared_id,
+        'descriptor': _recipe_descriptor_to_dict(descriptor),
+        'editVisitor': f'edit:{prepared_id}',
+        'editPreconditions': _get_preconditions(recipe, 'edit'),
+        'scanVisitor': f'scan:{prepared_id}' if is_scanning else None,
+        'scanPreconditions': _get_preconditions(recipe, 'scan') if is_scanning else [],
+    }
+
+    logger.info(f"PrepareRecipe response: {response}")
+    return response
+
+
+def _get_preconditions(recipe, phase: str) -> List[dict]:
+    """Get preconditions for a recipe phase.
+
+    For now, we add a type precondition to ensure only Python files are visited.
+    """
+    # Add precondition to only visit Python source files
+    return [{
+        'visitorName': 'org.openrewrite.rpc.internal.FindTreesOfType',
+        'visitorOptions': {'type': 'org.openrewrite.python.tree.Py'}
+    }]
+
+
+def handle_visit(params: dict) -> dict:
+    """Handle a Visit RPC request.
+
+    Applies a visitor to a tree and returns whether it was modified.
+
+    Args:
+        params: dict with 'visitor', 'sourceFileType', 'treeId', 'p' (context id), 'cursor'
+
+    Returns:
+        dict with 'modified' boolean
+    """
+    visitor_name = params.get('visitor')
+    source_file_type = params.get('sourceFileType')
+    tree_id = params.get('treeId')
+    p_id = params.get('p')
+    cursor_ids = params.get('cursor')
+
+    if visitor_name is None:
+        raise ValueError("'visitor' is required")
+    if tree_id is None:
+        raise ValueError("'treeId' is required")
+
+    logger.info(f"Visit: visitor={visitor_name}, treeId={tree_id}, p={p_id}")
+
+    # Get or create execution context
+    if p_id and p_id in _execution_contexts:
+        ctx = _execution_contexts[p_id]
+    else:
+        from rewrite import InMemoryExecutionContext
+        ctx = InMemoryExecutionContext()
+        if p_id:
+            _execution_contexts[p_id] = ctx
+
+    # Get the tree - fetch from Java if we don't have it locally
+    tree = local_objects.get(tree_id)
+    if tree is None:
+        tree = get_object_from_java(tree_id, source_file_type)
+
+    if tree is None:
+        raise ValueError(f"Tree not found: {tree_id}")
+
+    # Instantiate the visitor
+    visitor = _instantiate_visitor(visitor_name, ctx)
+
+    # Apply the visitor
+    from rewrite.visitor import Cursor
+    cursor = Cursor(None, Cursor.ROOT_VALUE)
+
+    before = tree
+    after = visitor.visit(tree, ctx, cursor)
+
+    # Update local objects with the result and determine if modified
+    # Use referential equality (identity comparison) to detect modifications
+    if after is None:
+        # Tree was deleted
+        if tree_id in local_objects:
+            del local_objects[tree_id]
+        modified = True
+    elif after is not before:
+        # Tree object changed - update both the tree_id entry and the new id entry
+        local_objects[tree_id] = after
+        if str(after.id) != tree_id:
+            local_objects[str(after.id)] = after
+        modified = True
+    else:
+        modified = False
+
+    logger.info(f"Visit result: modified={modified}, tree_id={tree_id}, before.id={before.id}, after.id={after.id if after else None}")
+    return {'modified': modified}
+
+
+def _instantiate_visitor(visitor_name: str, ctx):
+    """Instantiate a visitor from its name.
+
+    Visitor names can be:
+    - 'edit:<recipe_id>' - get the editor from a prepared recipe
+    - 'scan:<recipe_id>' - get the scanner from a prepared scanning recipe
+
+    For ScanningRecipes, the accumulator is persisted across calls so that
+    data collected during the scan phase is available during the edit and
+    generate phases.
+    """
+    if visitor_name.startswith('edit:'):
+        recipe_id = visitor_name[5:]
+        recipe = _prepared_recipes.get(recipe_id)
+        if recipe is None:
+            raise ValueError(f"Prepared recipe not found: {recipe_id}")
+
+        # Track phase transition
+        _recipe_phases[recipe_id] = 'edit'
+
+        # For ScanningRecipe, use the accumulated data from scan phase
+        from rewrite.recipe import ScanningRecipe
+        if isinstance(recipe, ScanningRecipe):
+            # Get existing accumulator or create new one
+            if recipe_id not in _recipe_accumulators:
+                _recipe_accumulators[recipe_id] = recipe.initial_value(ctx)
+            acc = _recipe_accumulators[recipe_id]
+            return recipe.editor_with_data(acc)
+
+        return recipe.editor()
+
+    elif visitor_name.startswith('scan:'):
+        recipe_id = visitor_name[5:]
+        recipe = _prepared_recipes.get(recipe_id)
+        if recipe is None:
+            raise ValueError(f"Prepared recipe not found: {recipe_id}")
+        from rewrite.recipe import ScanningRecipe
+        if not isinstance(recipe, ScanningRecipe):
+            raise ValueError(f"Recipe is not a scanning recipe: {recipe_id}")
+
+        # Check for phase transition (edit -> scan = new cycle)
+        # If we're transitioning from edit back to scan, clear the accumulator
+        if _recipe_phases.get(recipe_id) == 'edit':
+            _recipe_accumulators.pop(recipe_id, None)
+        _recipe_phases[recipe_id] = 'scan'
+
+        # Get existing accumulator or create new one
+        if recipe_id not in _recipe_accumulators:
+            _recipe_accumulators[recipe_id] = recipe.initial_value(ctx)
+        acc = _recipe_accumulators[recipe_id]
+
+        return recipe.scanner(acc)
+
+    else:
+        raise ValueError(f"Unknown visitor name format: {visitor_name}")
+
+
+def handle_generate(params: dict) -> dict:
+    """Handle a Generate RPC request.
+
+    Called by the recipe run cycle to generate new source files from scanning recipes.
+    For non-scanning recipes, returns an empty list.
+
+    The accumulator used here is the same one that was populated during the scan phase,
+    allowing recipes to generate files based on data collected across all source files.
+
+    Args:
+        params: dict with 'id' (prepared recipe id) and 'p' (context id)
+
+    Returns:
+        dict with 'ids' list and 'sourceFileTypes' list
+    """
+    recipe_id = params.get('id')
+    p_id = params.get('p')
+
+    if recipe_id is None:
+        raise ValueError("'id' is required")
+
+    logger.info(f"Generate: id={recipe_id}, p={p_id}")
+
+    recipe = _prepared_recipes.get(recipe_id)
+    if recipe is None:
+        raise ValueError(f"Prepared recipe not found: {recipe_id}")
+
+    # Get or create execution context
+    if p_id and p_id in _execution_contexts:
+        ctx = _execution_contexts[p_id]
+    else:
+        from rewrite import InMemoryExecutionContext
+        ctx = InMemoryExecutionContext()
+        if p_id:
+            _execution_contexts[p_id] = ctx
+
+    # Only scanning recipes can generate files
+    from rewrite.recipe import ScanningRecipe
+    if isinstance(recipe, ScanningRecipe):
+        # Use the persisted accumulator from the scan phase, or create new one if not available
+        if recipe_id in _recipe_accumulators:
+            acc = _recipe_accumulators[recipe_id]
+        else:
+            acc = recipe.initial_value(ctx)
+            _recipe_accumulators[recipe_id] = acc
+
+        generated = recipe.generate(acc, ctx)
+
+        ids = []
+        source_file_types = []
+        for sf in generated:
+            sf_id = str(sf.id)
+            local_objects[sf_id] = sf
+            ids.append(sf_id)
+            source_file_types.append(sf.__class__.__module__ + '.' + sf.__class__.__name__)
+
+        return {'ids': ids, 'sourceFileTypes': source_file_types}
+
+    # Non-scanning recipes don't generate files
+    return {'ids': [], 'sourceFileTypes': []}
+
+
 def handle_request(method: str, params: dict) -> Any:
     """Handle an RPC request."""
     handlers = {
@@ -531,6 +810,9 @@ def handle_request(method: str, params: dict) -> Any:
         'Print': handle_print,
         'Reset': handle_reset,
         'GetMarketplace': handle_get_marketplace,
+        'PrepareRecipe': handle_prepare_recipe,
+        'Visit': handle_visit,
+        'Generate': handle_generate,
     }
 
     handler = handlers.get(method)