Skip to content

Commit a2de58b

Browse files
committed
can fix tuple asserts
1 parent c0f1b24 commit a2de58b

2 files changed

Lines changed: 103 additions & 0 deletions

File tree

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import libcst as cst
2+
from typing import List
3+
from codemodder.codemods.api import BaseCodemod, ReviewGuidance
4+
from codemodder.codemods.utils_mixin import NameResolutionMixin
5+
6+
7+
class FixAssertTuple(BaseCodemod, NameResolutionMixin):
8+
NAME = "fix-assert-tuple"
9+
SUMMARY = "TODISimplify Boolean Expressions Using `startswith` and `endswith`"
10+
REVIEW_GUIDANCE = ReviewGuidance.MERGE_WITHOUT_REVIEW
11+
DESCRIPTION = "TODOUse tuple of matches instead of boolean expression"
12+
REFERENCES: list = []
13+
14+
def leave_SimpleStatementLine(
15+
self,
16+
original_node: cst.SimpleStatementLine,
17+
updated_node: cst.SimpleStatementLine,
18+
) -> cst.FlattenSentinel: # todo update
19+
if not self.filter_by_path_includes_or_excludes(
20+
self.node_position(original_node)
21+
):
22+
return updated_node
23+
24+
if len(updated_node.body) == 1 and isinstance(
25+
assert_node := updated_node.body[0], cst.Assert
26+
):
27+
match assert_test := assert_node.test:
28+
case cst.Tuple():
29+
if not assert_test.elements:
30+
return updated_node
31+
new_asserts = self._make_asserts(assert_node)
32+
# todo: report change for each added line
33+
return cst.FlattenSentinel(new_asserts)
34+
return updated_node
35+
36+
def _make_asserts(self, node: cst.Assert) -> List[cst.SimpleStatementLine]:
37+
return [
38+
cst.SimpleStatementLine(body=[cst.Assert(test=element.value, msg=node.msg)])
39+
for element in node.test.elements
40+
]
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
import pytest
2+
from tests.codemods.base_codemod_test import BaseCodemodTest
3+
from core_codemods.fix_assert_tuple import FixAssertTuple
4+
5+
6+
class TestFixAssertTuple(BaseCodemodTest):
7+
codemod = FixAssertTuple
8+
9+
def test_name(self):
10+
assert self.codemod.name() == "fix-assert-tuple"
11+
12+
@pytest.mark.parametrize(
13+
"input_code,expected_output",
14+
[
15+
(
16+
"""\
17+
print(1)
18+
assert (1, 2, )
19+
print(2)
20+
""",
21+
"""\
22+
print(1)
23+
assert 1
24+
assert 2
25+
print(2)
26+
""",
27+
),
28+
("""assert (1,)""", """assert 1"""),
29+
(
30+
"""assert ("one", Exception, [])""",
31+
"""\
32+
assert "one"
33+
assert Exception
34+
assert []""",
35+
),
36+
],
37+
)
38+
def test_change(self, tmpdir, input_code, expected_output):
39+
self.run_and_assert(tmpdir, input_code, expected_output)
40+
# assert len(self.file_context.codemod_changes) == 1
41+
42+
@pytest.mark.parametrize(
43+
"code",
44+
[
45+
"assert [1]",
46+
"assert 1",
47+
"assert ()",
48+
"assert (1)",
49+
"assert ('hi')",
50+
],
51+
)
52+
def test_no_change(self, tmpdir, code):
53+
self.run_and_assert(tmpdir, code, code)
54+
assert len(self.file_context.codemod_changes) == 0
55+
56+
def test_exclude_line(self, tmpdir):
57+
input_code = expected = """\
58+
assert (1, 2)
59+
"""
60+
lines_to_exclude = [1]
61+
self.assert_no_change_line_excluded(
62+
tmpdir, input_code, expected, lines_to_exclude
63+
)

0 commit comments

Comments
 (0)