Coverage for src / renaissance / syntax_tree / ast_rewriter.py: 86%
290 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-09-09 14:04 +0000
1import re
2import sys
3from collections.abc import Sequence
4from enum import Enum
5from typing import Protocol, Self, runtime_checkable
7from more_itertools import flatten
9from renaissance.common import Rewriter
10from renaissance.utils.text_utils import TextUtils
12from ..integrations.types import CompoundStatement
13from .ast_finder import ASTFinder
14from .match_finder import PatternMatch
17@runtime_checkable
18class Rewritable(Protocol):
19 offset: int
20 end_offset: int
21 extended_end_offset: int
22 filename: str
23 parent: Self
24 text: str
27class _RewriteActionType(Enum):
28 REPLACE = 1
29 INSERT_BEFORE = 2
30 INSERT_AFTER = 3
31 REMOVE = 4 # TODO: Why needed? Why isn't a REMOVE Action Type just a REPLACE Action Type (with an empty string)?
34DEFAULT_INDENT = 4
37class ASTRewriter:
38 def __init__(
39 self,
40 node,
41 encoding: str = sys.getfilesystemencoding(),
42 correct_indent: bool = True,
43 ) -> None:
44 self.__rewrites = _RewriteActions(node, encoding, correct_indent=correct_indent)
45 self.__filename = node.filename
47 def get_filename(self) -> str:
48 return self.__filename
50 def replace(
51 self,
52 new_content: str,
53 target: Rewritable | Sequence[Rewritable] | PatternMatch | Sequence[PatternMatch],
54 include_whitespace: bool = True,
55 include_comments: bool = True,
56 ):
57 self.__rewrites.add(
58 _RewriteActionType.REPLACE,
59 target,
60 new_content,
61 include_whitespace,
62 include_comments,
63 )
65 def remove(
66 self,
67 target: Rewritable | Sequence[Rewritable] | PatternMatch | Sequence[PatternMatch],
68 include_whitespace: bool = True,
69 include_comments: bool = True,
70 ):
71 self.__rewrites.add(_RewriteActionType.REMOVE, target, "", include_whitespace, include_comments)
73 def insert_before(
74 self,
75 new_content: str,
76 target: Rewritable | Sequence[Rewritable] | PatternMatch | Sequence[PatternMatch],
77 include_whitespace: bool = True,
78 include_comments: bool = True,
79 ):
80 self.__rewrites.add(
81 _RewriteActionType.INSERT_BEFORE,
82 target,
83 new_content,
84 include_whitespace,
85 include_comments,
86 )
88 def insert_after(
89 self,
90 new_content: str,
91 target: Rewritable | Sequence[Rewritable] | PatternMatch | Sequence[PatternMatch],
92 include_whitespace: bool = True,
93 include_comments: bool = True,
94 ):
95 self.__rewrites.add(
96 _RewriteActionType.INSERT_AFTER,
97 target,
98 new_content,
99 include_whitespace,
100 include_comments,
101 )
103 def apply_to_string(self) -> str:
104 return self.__rewrites.apply_to_string()
106 def apply(self) -> bytes:
107 if len(self.__rewrites.rewrites) == 0:
108 return self.__rewrites.content
109 return self.__rewrites.apply()
111 def has_changed(self) -> bool:
112 return len(self.__rewrites.rewrites) > 0
114 @staticmethod
115 def _get_comment_location(start_offset: int, stop_offset: int, content: bytes) -> tuple[int, int]:
116 return _RewriteActions.get_comment_location(start_offset, stop_offset, content)
119class _RewriteAction:
120 """Data container for a rewrite action to be applied later on to the AST."""
122 def __init__(
123 self,
124 action: _RewriteActionType,
125 target: Rewritable | Sequence[Rewritable] | PatternMatch | Sequence[PatternMatch],
126 replacement: str,
127 include_whitespace: bool,
128 include_comments: bool,
129 ) -> None:
130 self.action = action
131 self.target = target
132 self.replacement = replacement
133 self.nodes = self._get_nodes(target)
134 self.include_whitespace = include_whitespace
135 self.include_comments = include_comments
137 @staticmethod
138 def _get_nodes(
139 target: Rewritable | Sequence[Rewritable] | PatternMatch | Sequence[PatternMatch],
140 ) -> Sequence[Rewritable]:
141 if isinstance(target, Rewritable) or type(target).__name__ == "PythonASTNode":
142 return [target]
143 if isinstance(target, PatternMatch):
144 return target.nodes
145 assert isinstance(target, Sequence), "type of target violates its type requirements " + type(target).__name__
146 if len(target) > 0:
147 if isinstance(target[0], Rewritable): # TODO Why is part missing That is present on line 140, i.e.,
148 # or type(target).__name__ == "PythonASTNode"
149 return [n for n in target if isinstance(n, Rewritable)]
150 last = target[-1]
151 assert isinstance(last, PatternMatch), "type within Sequence violates its requirements " + type(last).__name__
152 return last.nodes
153 # TODO: is this correct? Can the other matches indeed be ignored?
154 return []
157class _RewriteActions:
158 """Data container for a list of rewrite actions to be applied later on to the AST."""
160 def __init__(
161 self,
162 node: Rewritable,
163 encoding: str,
164 correct_indent: bool,
165 rewrites: list[_RewriteAction] | None = None,
166 ) -> None:
167 self.rewrites: list[_RewriteAction] = rewrites or []
168 self.node = node
169 self.encoding = encoding
170 # self.content = self.node.root.binary_file_content()[self.node.offset : self.node.extended_end_offset]
171 self.content = node.text.encode(sys.getfilesystemencoding())
172 self.correct_indent = correct_indent
174 def add(
175 self,
176 action: _RewriteActionType,
177 target: Rewritable | Sequence[Rewritable] | PatternMatch | Sequence[PatternMatch],
178 replacement: str,
179 include_whitespace: bool,
180 include_comments: bool,
181 ):
182 rewrite = _RewriteAction(action, target, replacement, include_whitespace, include_comments)
183 self.add_rewrite(rewrite)
185 def add_rewrite(self, rewrite: _RewriteAction):
186 self.rewrites.append(rewrite)
188 def apply(self) -> bytes:
189 rewriter = Rewriter(self.content[:])
191 for rewrite in self.rewrites:
192 # skip nested rewrites as they are handled recursively by the parent rewrite
193 # except for if the rewrite node is the root node
194 if any(self.__is_ancestor_in_nodes(n) for n in rewrite.nodes if n != self.node):
195 continue
196 new_content, nodelist = self.__prepare_replacement_content(rewrite.replacement, rewrite.target)
197 if rewrite.action == _RewriteActionType.REPLACE:
198 self.__replace(
199 rewriter,
200 new_content,
201 nodelist,
202 rewrite.include_whitespace,
203 rewrite.include_comments,
204 )
205 elif rewrite.action == _RewriteActionType.INSERT_BEFORE:
206 self.__insert(
207 rewriter,
208 new_content,
209 True,
210 nodelist,
211 rewrite.include_whitespace,
212 rewrite.include_comments,
213 )
214 elif rewrite.action == _RewriteActionType.INSERT_AFTER:
215 self.__insert(
216 rewriter,
217 new_content,
218 False,
219 nodelist,
220 rewrite.include_whitespace,
221 rewrite.include_comments,
222 )
223 elif rewrite.action == _RewriteActionType.REMOVE:
224 self.__remove(
225 rewriter,
226 nodelist,
227 rewrite.include_whitespace,
228 rewrite.include_comments,
229 )
230 return rewriter.apply()
232 def apply_to_string(self) -> str:
233 return self.apply().decode(self.encoding)
235 def __is_ancestor_in_nodes(self, node: Rewritable) -> bool:
236 """Check if the given node is a descendant of any nodes in the rewrite list.
238 Args:
239 node (Rewritable): The node to check.
241 Returns:
242 bool: True if the node is a descendant of any nodes in the rewrite list, False otherwise.
244 """
245 rewrite_nodes = list(flatten(rewrite.nodes for rewrite in self.rewrites))
247 # need to test
248 # 1
249 # | node |
250 # |rew|
251 # 2
252 # | rew |
253 # |node|
254 def no_conflict(node1, rew):
255 return not (node1.end_offset < rew.offset or node1.offset > rew.end_offset)
257 result = any(no_conflict(node, rew) for rew in rewrite_nodes)
259 return result and False # TODO: Why `and False`
261 def __replace(
262 self,
263 rewriter: Rewriter,
264 new_content: str,
265 nodes: Sequence[Rewritable],
266 include_whitespace: bool,
267 include_comments: bool,
268 ):
269 """Replaces the content of the given node(s) with new content.
271 Args:
272 nodes (Sequence[Rewritable]): The nodes whose content is to be replaced.
273 new_content (str): The new content to insert in the specified range.
275 """
276 if not nodes:
277 return
278 start_offset, end_offset = _RewriteActions.__correct_for_comments_and_whitespace(
279 self.node.offset,
280 self.content,
281 include_whitespace,
282 include_comments,
283 nodes,
284 )
285 # start_offset =nodes[0].get_start_offset()
286 # end_offset =nodes[-1].get_start_offset()+nodes[-1].get_length()+1
287 indent = self.derive_indent(start_offset)
288 if self.correct_indent:
289 if new_content.startswith("\n"):
290 # a blank first line would otherwise strand the original indent as trailing whitespace
291 start_offset -= indent
292 new_content = TextUtils.shift_right(new_content, indent, start_line=1)
293 self.__replace_bytes(rewriter, start_offset, end_offset, new_content)
295 def __remove(
296 self,
297 rewriter: Rewriter,
298 nodes: Sequence[Rewritable],
299 include_whitespace: bool = False,
300 include_comments: bool = False,
301 ):
302 """Removes a list of AST nodes from the content, optionally including surrounding whitespace and comments.
304 Args:
305 nodes (Sequence[Rewritable]): The list of AST nodes to remove.
306 include_whitespace (bool, optional): Whether to include surrounding whitespace in the removal. Defaults to False.
307 include_comments (bool, optional): Whether to include surrounding comments in the removal. Defaults to False.
309 Returns:
310 None
312 """
313 if not nodes:
314 return
316 start_offset, end_offset = _RewriteActions.__correct_for_comments_and_whitespace(
317 self.node.offset,
318 self.content,
319 include_whitespace,
320 include_comments,
321 nodes,
322 )
323 indent = self.derive_indent(start_offset)
324 # remove the indent in front of it
325 start_offset -= indent
326 # remove the line if it is empty
327 if start_offset > 0 and self.content[start_offset - 1] == ord("\n") and self.content[end_offset] == ord("\n"):
328 start_offset -= 1
329 self.__replace_bytes(rewriter, start_offset, end_offset, "")
331 def derive_indent(self, start_offset: int) -> int:
332 indent = 0 # len(nodes[0].indent)
333 if start_offset > 0:
334 while len(self.content) > (start_offset - indent - 1) and self.content[start_offset - indent - 1] in [32]:
335 indent += 1
336 return indent
338 def __insert(
339 self,
340 rewriter: Rewriter,
341 new_content: str,
342 before: bool,
343 nodes: Sequence[Rewritable],
344 include_whitespace: bool,
345 include_comments: bool,
346 ):
347 if not nodes:
348 return
349 content = self.content
350 indent = TextUtils.get_spaces_before(content, nodes[0].offset)
351 spaces = " " * indent
352 # if flattened_nodes[-1] has a new line after white space then we need to add a new line:
353 ext_start_offset, ext_end_offset = _RewriteActions.__correct_for_comments_and_whitespace(
354 self.node.offset,
355 self.content,
356 include_whitespace,
357 include_comments,
358 nodes,
359 )
360 white_space = (
361 ""
362 if not include_whitespace
363 else "\n" + spaces
364 if ext_end_offset < len(content) and content[ext_end_offset] in b"\n"
365 else spaces
366 )
367 # indent the new content except the first line
368 new_content = TextUtils.shift_right(new_content, indent, start_line=1)
370 if before:
371 # restore the node's original indent, consumed as new_content's first-line indent
372 if not white_space and new_content.endswith("\n"):
373 new_content += spaces
374 self.__replace_bytes(rewriter, ext_start_offset, ext_start_offset, new_content + white_space)
375 else:
376 self.__replace_bytes(rewriter, ext_end_offset, ext_end_offset, white_space + new_content)
378 def __replace_bytes(self, rewriter: Rewriter, start: int, end: int, new_content: str) -> None:
379 """Replaces the content in the specified range with new content.
381 Args:
382 start (int): The starting index of the range to be replaced.
383 end (int): The ending index of the range to be replaced.
384 new_content (str): The new content to insert in the specified range.
386 """
387 rewriter.replace(start, end, new_content.encode(self.encoding))
389 def __compose_replacement(self, replacement: str, matches: Sequence[PatternMatch]) -> str:
390 all_placeholders = {p: n for m in matches for p, n in m.expansions.items()}
391 for placeholder, nodes in all_placeholders.items():
392 quoted_placeholder = re.escape(placeholder)
393 raw_signature = self.__get_texts(nodes)
394 # replacement = replacement.replace(placeholder, raw_signature)
395 while placeholder in replacement:
396 pattern = re.compile(r"( *)" + quoted_placeholder)
397 matcher = pattern.search(replacement)
399 if matcher:
400 spaces = matcher[1]
401 place_holder_length = len(placeholder)
402 index = replacement.index(placeholder)
403 # TODO a regex may be provided between backticks and the groups are used. This needs a better design
404 # A preferable solution is to pass a transformer function to the compose_replacement
405 if index + place_holder_length < len(replacement) and replacement[index + place_holder_length] == "`":
406 # ` ` means get regex
407 end_index = replacement.index("`", index + place_holder_length + 1)
408 if not end_index:
409 raise ValueError("No closing ` found")
410 regex = replacement[index + place_holder_length + 1 : end_index]
411 regex_match = re.match(regex, raw_signature)
412 if regex_match:
413 raw_signature = "".join(regex_match.groups())
414 place_holder_length = end_index - index + 1
415 indent_replacement = raw_signature.replace("\n", "\n" + spaces)
416 if (
417 placeholder.startswith("$$")
418 and index + place_holder_length < len(replacement)
419 and replacement[index + place_holder_length] == ";"
420 ):
421 place_holder_length += 1
422 # replace the placeholder with the indent replacement
423 replacement = replacement[:index] + indent_replacement + replacement[index + place_holder_length :]
424 else:
425 print("Match doesn't match unexpectedly")
426 return replacement
428 def __get_texts(self, nodes: Sequence[Rewritable]) -> str:
429 if len(nodes) == 1:
430 return self.__get_text(nodes[0])
431 # Use a ASTRewriter to only rewrite exactly that what needs to be rewritten
432 rewriter = ASTRewriter(nodes[0], self.encoding, correct_indent=False)
433 for node in nodes:
434 rs = self.__get_text(node)
435 org_rs = node.text
436 if rs != org_rs:
437 rewriter.replace(rs, node)
438 result = rewriter.apply_to_string()
439 indent = self.derive_indent(nodes[0].offset)
440 return TextUtils.shift_left(result, indent, start_line=1)
442 def __get_text(self, node: Rewritable) -> str:
443 if self._should_skip(node):
444 return ""
446 if node == self.node:
447 return node.text
448 # the descendants may need to be rewritten as well
449 # rewrites = [rewrite for rewrite in self.rewrites if any(node.is_ancestor_of(rewrite_node)
450 # for rewrite_node in rewrite.nodes)]
451 rewrites = [rewrite for rewrite in self.rewrites if any(node.is_ancestor_of(rewrite_node) for rewrite_node in rewrite.nodes)]
452 if rewrites:
453 rewriter = _RewriteActions(node, self.encoding, self.correct_indent, rewrites)
454 return rewriter.apply_to_string()
455 return node.text
457 def __prepare_replacement_content(
458 self,
459 new_content: str,
460 target: PatternMatch | Rewritable | Sequence[Rewritable],
461 ) -> tuple[str, Sequence[Rewritable]]:
462 if isinstance(target, PatternMatch):
463 new_content = self.__compose_replacement(new_content, [target])
464 node_list = target.nodes
465 else:
466 node_list = (
467 [target] if (isinstance(target, Rewritable) or type(target).__name__ == "PythonASTNode") else target
468 ) # TODO How to make a Sequence[Rewritable] as type hints also show list[Rewritable]?
469 return new_content, node_list
471 def _should_skip(self, node: Rewritable):
472 """If the node is not the first node of a pattern match it should be skipped."""
473 return any(node in rewrite.nodes[1:] for rewrite in self.rewrites if isinstance(rewrite.target, PatternMatch))
475 @staticmethod
476 def _get_parent_statement(node: Rewritable):
477 parent = node
478 while parent and not parent.is_statement:
479 parent = parent.parent
480 return parent
482 @staticmethod
483 def __correct_for_comments_and_whitespace(
484 offset: int,
485 content: bytes,
486 include_whitespace: bool,
487 include_comments: bool,
488 nodes: Sequence[Rewritable],
489 ):
490 start_offset = nodes[0].offset - offset
491 end_offset = nodes[-1].extended_end_offset - offset
492 if include_comments:
493 preceding_node = nodes[0].preceding_sibling
494 parent = nodes[0].parent
495 start_comment_location = 0
496 if preceding_node:
497 # start after the comment of the preceding node
498 start_comment_location = preceding_node.extended_end_offset - offset
499 preceding_end_offset = _RewriteActions.__get_comment_after_location(start_comment_location, start_offset, content)
500 if preceding_end_offset != (-1, -1):
501 start_comment_location = preceding_end_offset[1]
502 elif parent:
503 start_comment_location = parent.offset - offset
504 # get the comment belonging to the preceding node
505 extended_location = _RewriteActions.get_comment_location(start_comment_location, start_offset, content)
506 if extended_location != (-1, -1):
507 start_offset = extended_location[0]
508 next_sibling = nodes[-1].next_sibling
509 end_comment_location = next_sibling.offset - offset if next_sibling else parent.end_offset - offset if parent else len(content)
510 location_after_comment = _RewriteActions.__get_comment_after_location(end_offset, end_comment_location, content)
511 if location_after_comment != (-1, -1):
512 end_offset = location_after_comment[1]
513 if include_whitespace:
514 end_offset = _RewriteActions.__extend_with_whitespace(end_offset, content)
515 return start_offset, end_offset
517 def cor_offset(self, offset: int):
518 return offset - self.node.offset
520 @staticmethod
521 def get_comment_location(start_offset: int, stop_offset: int, content: bytes) -> tuple[int, int]:
522 """Get the location of the comment before the location, but after the stop_location
523 a comment is a line that starts with // or a block that starts with /* and ends with */
524 or a line that starts with #.
525 """
526 # search last occurrence of //, /*, # in a byte array
527 comment_start = content.rfind(b"//", start_offset, stop_offset)
528 if comment_start != -1:
529 comment_end = _RewriteActions.__get_end_of_line(content, comment_start)
530 return comment_start, comment_end
531 comment_start = content.rfind(b"/*", start_offset, stop_offset)
532 if comment_start != -1:
533 comment_end = content.find(b"*/", comment_start, stop_offset)
534 if comment_end != -1:
535 comment_end += len("*/")
536 return comment_start, comment_end
537 comment_start = content.rfind(b"#", start_offset, stop_offset)
538 if comment_start != -1:
539 comment_end = _RewriteActions.__get_end_of_line(content, comment_start)
540 return comment_start, comment_end
541 return -1, -1
543 @staticmethod
544 def __extend_with_whitespace(start_offset: int, content: bytes) -> int:
545 end_location = _RewriteActions.__get_end_of_line(content, start_offset)
546 text = content[start_offset:end_location]
547 for byt in text:
548 if byt not in b" \t":
549 return start_offset
550 return end_location
552 @staticmethod
553 def __get_comment_after_location(start_offset: int, end_offset: int, content: bytes) -> tuple[int, int]:
554 """Get the location of the comment before the location, but after the stop_location
555 a comment is a line that starts with // or a block that starts with /* and ends with */
556 or a line that starts with #.
557 """
558 line_end_offset = _RewriteActions.__get_end_of_line(content, start_offset)
559 if line_end_offset == -1:
560 line_end_offset = len(content)
561 comment_start = content.find(b"//", start_offset, line_end_offset)
562 if comment_start == -1:
563 comment_start = content.rfind(b"#", start_offset, line_end_offset)
564 if comment_start != -1:
565 return comment_start, line_end_offset
566 comment_start = content.rfind(b"/*", start_offset, line_end_offset)
567 if comment_start != -1:
568 # a block comment must start on the same line but doesn't have to finish on the same line
569 comment_end = content.find(b"*/", comment_start, end_offset)
570 if comment_end != -1:
571 comment_end += len("*/")
572 return comment_start, comment_end
573 return -1, -1
575 @staticmethod
576 def __get_end_of_line(content: bytes, start: int):
577 location = content.find(b"\n", start)
578 if location == -1:
579 return len(content)
580 return location
582 @staticmethod
583 def __get_depth(node: Rewritable) -> int:
584 depth = 0
585 parent = node.parent
586 while parent:
587 if ASTFinder.matches_kind(parent, CompoundStatement):
588 depth += 1
589 parent = parent.parent
590 return depth