Coverage for src/renaissance/syntax_tree/recipe_ast_processor.py: 100%

1"""AI: Decorator support for annotating recipe actions on AST processors.""" 

2 

3import functools 

4from collections.abc import Callable, Sequence 

5from typing import Any, TypeVar 

6 

7from .ast_processor import ASTProcessor 

8from .batch_ast_processor import BatchASTProcessor, IterableProvider 

9 

10TRecipe = TypeVar("TRecipe") 

11TFunc = Callable[..., Any] 

12 

13 

14def annotate_decorator(foreign_decorator: TFunc, name: str): 

15 """AI: Wrap foreign_decorator so the decorated function also records its own decorator and recipe-action name.""" 

16 

17 def new_decorator(func: TFunc) -> TFunc: 

18 r = foreign_decorator(func) # apply foreignDecorator, like call to foreignDecorator(method) would have done 

19 r.decorator = new_decorator # keep track of decorator 

20 r.recipe_action = name 

21 return r 

22 

23 new_decorator.__name__ = foreign_decorator.__name__ 

24 new_decorator.__doc__ = foreign_decorator.__doc__ 

25 return new_decorator 

26 

27 

28def get_methods_with_decorator(cls: Any, decorator: TFunc): 

29 """AI: Yield all methods of cls that were annotated with the given recipe-action decorator.""" 

30 for maybe_decorated in cls.__dict__.values(): 

31 if hasattr(maybe_decorated, "recipe_action") and maybe_decorated.recipe_action == decorator.__name__: 

32 yield maybe_decorated 

33 

34 

35# Decorators 

36 

37 

38def final_action() -> TFunc: 

39 """AI: Decorator marking a recipe method as a final action, run once after all steps complete.""" 

40 

41 def final_action_decorator(func: TFunc) -> TFunc: 

42 @functools.wraps(func) 

43 def final_action_wrapper(recipe: TFunc): 

44 func(recipe) 

45 

46 return final_action_wrapper 

47 

48 return annotate_decorator(final_action_decorator, final_action.__name__) 

49 

50 

51def recipe_step(order: int = 0, repeat: bool = False) -> TFunc: 

52 """AI: Decorator marking a recipe method as a numbered step, optionally repeatable, in the recipe's workflow.""" 

53 

54 def recipe_step_decorator(func: TFunc) -> TFunc: 

55 @functools.wraps(func) 

56 def recipe_step_wrapper(step: int, recipe: TFunc, ast_processor: ASTProcessor): 

57 if step == order and (repeat or ast_processor.repeat_step == 0): 

58 result = func(recipe, ast_processor) 

59 

60 def callable_result(): 

61 if result: 

62 result() 

63 return func.__name__ 

64 

65 return callable_result() 

66 return None 

67 

68 return recipe_step_wrapper 

69 

70 return annotate_decorator(recipe_step_decorator, recipe_step.__name__) 

71 

72 

73def after_step(step: str) -> TFunc: 

74 """AI: Decorator marking a recipe method to run only after the named step has already executed.""" 

75 

76 def after_step_decorator(func: TFunc) -> TFunc: 

77 @functools.wraps(func) 

78 def after_step_wrapper(preceding_methods: Sequence[str], recipe: TFunc): 

79 if step in preceding_methods: 

80 func(recipe) 

81 

82 return after_step_wrapper 

83 

84 return annotate_decorator(after_step_decorator, after_step.__name__) 

85 

86 

87class RecipeASTProcessor[TRecipe]: 

88 """AI: Processor that runs a recipe's steps across a batch of files, optionally in parallel.""" 

89 

90 def __init__( 

91 self, 

92 recipe: TRecipe, 

93 iterable_provider: IterableProvider, 

94 file_filter: str, 

95 in_memory: bool = False, 

96 max_processes: int = 4, 

97 ): 

98 """AI: Prepare a processor that runs a recipe's steps across a batch of files, optionally in parallel.""" 

99 self.__recipe: TRecipe = recipe 

100 self.__batch_processor = BatchASTProcessor(in_memory=in_memory, max_processes=max_processes) 

101 self.__iterableProvider = iterable_provider 

102 self.__file_filter = file_filter 

103 

104 def run(self): 

105 """AI: Run each of the recipe's decorated steps across the batch of files and return the collected results.""" 

106 actions: list[TFunc] = [] 

107 results: list[Any] = [] 

108 for idx, recipe_step_method in enumerate(get_methods_with_decorator(type(self.__recipe), recipe_step)): 

109 results.append(None) 

110 

111 def recipe_action(ast_processor: ASTProcessor): 

112 result = recipe_step_method(step, self.__recipe, ast_processor) 

113 if result: 

114 results[idx] = result 

115 

116 actions.append(recipe_action) 

117 

118 after_step_actions: list[TFunc] = [] 

119 for after_step_method in get_methods_with_decorator(self.__recipe.__class__, after_step): 

120 

121 def after_step_action(): 

122 after_step_method(results, self.__recipe) 

123 

124 after_step_actions.append(after_step_action) 

125 

126 step = 0 

127 while len(actions) > 0: 

128 for idx in range(len(results)): 

129 results[idx] = None 

130 self.__batch_processor.repeat(self.__iterableProvider, actions, self.__file_filter) 

131 if all(result is None for result in results): 

132 break 

133 for after_step_action in after_step_actions: 

134 after_step_action() 

135 step += 1 

136 

137 for method in get_methods_with_decorator(self.__recipe.__class__, final_action): 

138 method(self.__recipe)