forked from pixee/codemodder-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathfix_assert_tuple.py
More file actions
55 lines (49 loc) · 2.12 KB
/
fix_assert_tuple.py
File metadata and controls
55 lines (49 loc) · 2.12 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import libcst as cst
from typing import List, Union
from core_codemods.api import Metadata, ReviewGuidance, SimpleCodemod
from codemodder.change import Change
from codemodder.codemods.utils_mixin import NameResolutionMixin
class FixAssertTuple(SimpleCodemod, NameResolutionMixin):
metadata = Metadata(
name="fix-assert-tuple",
summary="Fix `assert` on Non-Empty Tuple Literal",
review_guidance=ReviewGuidance.MERGE_AFTER_CURSORY_REVIEW,
references=[],
)
change_description = "Separate assertion on a non-empty tuple literal into multiple assert statements."
def leave_SimpleStatementLine(
self,
original_node: cst.SimpleStatementLine,
updated_node: cst.SimpleStatementLine,
) -> Union[cst.FlattenSentinel, cst.SimpleStatementLine]:
if not self.filter_by_path_includes_or_excludes(
self.node_position(original_node)
):
return updated_node
if len(updated_node.body) == 1 and isinstance(
assert_node := updated_node.body[0], cst.Assert
):
match assert_test := assert_node.test:
case cst.Tuple():
if not assert_test.elements:
return updated_node
new_asserts = self._make_asserts(assert_node)
self._report_new_lines(original_node, len(new_asserts))
return cst.FlattenSentinel(new_asserts)
return updated_node
def _make_asserts(self, node: cst.Assert) -> List[cst.SimpleStatementLine]:
return [
cst.SimpleStatementLine(body=[cst.Assert(test=element.value, msg=node.msg)])
for element in node.test.elements
]
def _report_new_lines(
self, original_node: cst.SimpleStatementLine, newlines_count: int
):
start_line = self.node_position(original_node).start.line
for idx in range(newlines_count):
self.file_context.codemod_changes.append(
Change(
lineNumber=start_line + idx,
description=self.change_description,
)
)