Coverage for src/renaissance/syntax_tree/batch_ast_processor.py: 98%

1"""AI: Process multiple ASTs in parallel by applying a shared action to each.""" 

2 

3import concurrent.futures 

4import re 

5from collections.abc import Callable, Iterable, Sequence 

6from functools import partial 

7from typing import Any 

8 

9from .ast_factory import ASTFactory 

10from .ast_node import ASTNode 

11from .ast_processor import ASTProcessor 

12 

13AST_FACTORY_AND_ATU = tuple[ASTFactory, ASTNode] 

14Action = Callable[[ASTProcessor], Callable[[], Any] | None] 

15IterableProvider = Callable[[], Iterable[AST_FACTORY_AND_ATU]] 

16 

17 

18class BatchASTProcessor: 

19 """AI: Process multiple ASTs in parallel by applying a shared action to each.""" 

20 

21 def __init__(self, in_memory: bool = False, max_processes: int = 4): 

22 """Initialize the BatchASTProcessor. 

23 

24 Args: 

25 in_memory (bool): Flag to indicate if processing should be done in memory. Defaults to False. 

26 max_processes (int): The maximum number of processes to use. Defaults to 4. 

27 

28 """ 

29 self.in_memory: bool = in_memory 

30 self.in_memory_files: dict[str, str] = {} 

31 self.max_processes = max_processes 

32 

33 def once( 

34 self, 

35 iterable: Iterable[AST_FACTORY_AND_ATU] | IterableProvider, 

36 actions: Action | Sequence[Action], 

37 file_filter: str | re.Pattern[str] | None = None, 

38 ) -> None: 

39 """Process a given iterable of ATU objects or an IterableProvider with specified actions. 

40 

41 Args: 

42 iterable (Iterable[ATU] | IterableProvider): The iterable or provider of ATU objects to process. 

43 actions (Action | Sequence[Action]): The action or sequence of actions to apply to each item in the iterable. 

44 file_filter (Optional[str | re.Pattern], optional): A filter to apply to file names. Defaults to None. 

45 

46 Returns: 

47 bool: True if processing was successful, False otherwise. 

48 

49 """ 

50 iterable = iterable() if callable(iterable) else iterable 

51 self.__process(iterable, actions, self.in_memory, file_filter) 

52 

53 def repeat( 

54 self, 

55 iterable_provider: IterableProvider, 

56 actions: Action | Sequence[Action], 

57 file_filter: str | re.Pattern[str] | None = None, 

58 max_repeat: int = 5, 

59 ) -> None: 

60 """Repeats the processing of items provided by the iterableProvider until no changes left. 

61 

62 Up to a maximum number of times. 

63 

64 Args: 

65 iterable_provider (IterableProvider): A provider that yields items to be processed. 

66 actions (Action | Sequence[Action]): A single action or a sequence of actions to be performed on each item. 

67 file_filter (Optional[str | re.Pattern], optional): A filter to apply to the files being processed. Defaults to None. 

68 max_repeat (int, optional): The maximum number of times to repeat the processing. Defaults to 5. 

69 

70 Returns: 

71 bool: True if the processing still yields changes, False otherwise. 

72 

73 """ 

74 self.__process(iterable_provider(), actions, self.in_memory, file_filter, max_repeat) 

75 

76 def __process( 

77 self, 

78 iterable: Iterable[tuple[ASTFactory, ASTNode]], 

79 actions: Action | Sequence[Action], 

80 in_memory: bool = False, 

81 file_filter: str | re.Pattern[str] | None = None, 

82 max_repeat: int = 1, 

83 ) -> None: 

84 filter_pattern = ( 

85 file_filter if isinstance(file_filter, re.Pattern) else re.compile(file_filter) if file_filter is not None else None 

86 ) 

87 

88 def is_eligible(item: tuple[ASTFactory, ASTNode]) -> bool: 

89 return BatchASTProcessor.__eligible_file(filter_pattern, item) 

90 

91 actions = actions if isinstance(actions, Sequence) else [actions] 

92 # use parallel processing possible here 

93 partial_process_item = partial( 

94 process_atu, 

95 self=self, 

96 actions=actions, 

97 in_memory=in_memory, 

98 max_repeat=max_repeat, 

99 ) 

100 with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_processes) as executor: 

101 for results in executor.map(partial_process_item, filter(is_eligible, iterable)): 

102 for my_callable in results: 

103 # the post-processing is done in the main thread 

104 my_callable() 

105 

106 def _replace_if_in_memory(self, item: AST_FACTORY_AND_ATU) -> AST_FACTORY_AND_ATU: 

107 if self.in_memory and self.in_memory_files.get(item[1].filename): 

108 return item[0], item[0].create_from_text( 

109 self.in_memory_files[item[1].filename], 

110 item[1].filename, 

111 ) 

112 return item 

113 

114 @staticmethod 

115 def __eligible_file(file_filter: re.Pattern[str] | None, item: AST_FACTORY_AND_ATU) -> bool: 

116 return file_filter is None or file_filter.match(item[1].filename) is not None 

117 

118 

119def process_atu( 

120 atu: AST_FACTORY_AND_ATU, 

121 self: BatchASTProcessor, 

122 actions: Sequence[Action], 

123 in_memory: bool, 

124 max_repeat: int, 

125) -> Sequence[Callable[[], None]]: 

126 """AI: Run the given actions against one ATU, repeating up to max_repeat times, and return deferred result callables.""" 

127 atu = self._replace_if_in_memory(atu) 

128 ast_processor = ASTProcessor(atu[1], atu[0], in_memory) 

129 results: list[Callable[[], None]] = [] 

130 

131 for repeat in range(max_repeat): 

132 for action in actions: 

133 ast_processor.repeat_step = repeat 

134 result = action(ast_processor) 

135 if result: 

136 results.append(result) 

137 has_changed = ast_processor.has_changed() 

138 if not has_changed: 

139 return results 

140 ast_processor = ast_processor.commit() 

141 if self.in_memory: 

142 self.in_memory_files[ast_processor.filename] = ast_processor.apply_to_string() 

143 return results