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

1import concurrent.futures 

2import re 

3from collections.abc import Callable, Iterable, Sequence 

4from functools import partial 

5from typing import Any 

6 

7from .ast_factory import ASTFactory 

8from .ast_node import ASTNode 

9from .ast_processor import ASTProcessor 

10 

11AST_FACTORY_AND_ATU = tuple[ASTFactory, ASTNode] 

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

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

14 

15 

16class BatchASTProcessor: 

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

18 """Initialize the BatchASTProcessor. 

19 

20 Args: 

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

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

23 

24 """ 

25 self.in_memory: bool = in_memory 

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

27 self.max_processes = max_processes 

28 

29 def once( 

30 self, 

31 iterable: Iterable[AST_FACTORY_AND_ATU] | IterableProvider, 

32 actions: Action | Sequence[Action], 

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

34 ) -> None: 

35 """Processes a given iterable of ATU objects or an IterableProvider with specified actions. 

36 

37 Args: 

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

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

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

41 

42 Returns: 

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

44 

45 """ 

46 iterable = iterable() if callable(iterable) else iterable 

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

48 

49 def repeat( 

50 self, 

51 iterable_provider: IterableProvider, 

52 actions: Action | Sequence[Action], 

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

54 max_repeat: int = 5, 

55 ) -> None: 

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

57 Up to a maximum number of times. 

58 

59 Args: 

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

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

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

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

64 

65 Returns: 

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

67 

68 """ 

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

70 

71 def __process( 

72 self, 

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

74 actions: Action | Sequence[Action], 

75 in_memory: bool = False, 

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

77 max_repeat: int = 1, 

78 ) -> None: 

79 filter_pattern = ( 

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

81 ) 

82 

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

84 return BatchASTProcessor.__eligible_file(filter_pattern, item) 

85 

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

87 # use parallel processing possible here 

88 partial_process_item = partial( 

89 process_atu, 

90 self=self, 

91 actions=actions, 

92 in_memory=in_memory, 

93 max_repeat=max_repeat, 

94 ) 

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

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

97 for my_callable in results: 

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

99 my_callable() 

100 

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

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

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

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

105 item[1].filename, 

106 ) 

107 return item 

108 

109 @staticmethod 

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

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

112 

113 

114def process_atu( 

115 atu: AST_FACTORY_AND_ATU, 

116 self: BatchASTProcessor, 

117 actions: Sequence[Action], 

118 in_memory: bool, 

119 max_repeat: int, 

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

121 atu = self._replace_if_in_memory(atu) 

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

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

124 

125 for repeat in range(max_repeat): 

126 for action in actions: 

127 ast_processor.repeat_step = repeat 

128 result = action(ast_processor) 

129 if result: 

130 results.append(result) 

131 has_changed = ast_processor.has_changed() 

132 if not has_changed: 

133 return results 

134 ast_processor = ast_processor.commit() 

135 if self.in_memory: 

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

137 return results