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
67 lines (56 loc) · 2.37 KB
/
fix_assert_tuple.py
File metadata and controls
67 lines (56 loc) · 2.37 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
56
57
58
59
60
61
62
63
64
65
66
67
from typing import List, Union
import libcst as cst
from codemodder.codemods.libcst_transformer import (
LibcstResultTransformer,
LibcstTransformerPipeline,
)
from codemodder.codemods.utils_mixin import NameResolutionMixin
from codemodder.codetf import Change
from core_codemods.api import Metadata, ReviewGuidance
from core_codemods.api.core_codemod import CoreCodemod
class FixAssertTupleTransform(LibcstResultTransformer, NameResolutionMixin):
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 len(original_node.body) == 1 and isinstance(
assert_node := original_node.body[0], cst.Assert
):
match assert_test := assert_node.test:
case cst.Tuple():
if not self.node_is_selected(assert_test):
return updated_node
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,
)
)
FixAssertTuple = CoreCodemod(
metadata=Metadata(
name="fix-assert-tuple",
summary="Fix `assert` on Non-Empty Tuple Literal",
review_guidance=ReviewGuidance.MERGE_AFTER_CURSORY_REVIEW,
references=[],
),
transformer=LibcstTransformerPipeline(FixAssertTupleTransform),
detector=None,
)