Codemod fix-assert-tuple

integration_tests/test_fix_assert_tuple.py

@@ -0,0 +1,17 @@
+from core_codemods.fix_assert_tuple import FixAssertTuple
+from integration_tests.base_test import (
+    BaseIntegrationTest,
+    original_and_expected_from_code_path,
+)
+
+
+class TestFixAssertTuple(BaseIntegrationTest):
+    codemod = FixAssertTuple
+    code_path = "tests/samples/fix_assert_tuple.py"
+    original_code, expected_new_code = original_and_expected_from_code_path(
+        code_path, [(0, "assert 1 == 1\n"), (1, "assert 2 == 2\n")]
+    )
+    expected_diff = "--- \n+++ \n@@ -1 +1,2 @@\n-assert (1 == 1, 2 == 2)\n+assert 1 == 1\n+assert 2 == 2\n"
+    expected_line_change = "1"
+    change_description = FixAssertTuple.change_description
+    num_changes = 2

src/codemodder/scripts/generate_docs.py

@@ -218,6 +218,10 @@ class DocMetadata:
         importance="Low",
         guidance_explained="We believe this change is safe and will not cause any issues.",
     ),
+    "fix-assert-tuple": DocMetadata(
+        importance="Medium",
+        guidance_explained="An `assert` statement on a non-empty tuple is likely unintended and should be rewritten. However, the new change may result in assertion failures that should be reviewed.",
+    ),
 }
 
 

src/core_codemods/__init__.py

@@ -48,6 +48,7 @@
 from .replace_flask_send_file import ReplaceFlaskSendFile
 from .fix_empty_sequence_comparison import FixEmptySequenceComparison
 from .remove_assertion_in_pytest_raises import RemoveAssertionInPytestRaises
+from .fix_assert_tuple import FixAssertTuple
 
 registry = CodemodCollection(
     origin="pixee",
@@ -102,5 +103,6 @@
         ReplaceFlaskSendFile,
         FixEmptySequenceComparison,
         RemoveAssertionInPytestRaises,
+        FixAssertTuple,
     ],
 )

src/core_codemods/docs/pixee_python_fix-assert-tuple.md

@@ -0,0 +1,9 @@
+An assertion on a non-empty tuple will always evaluate to `True`. This means that `assert` statements involving non-empty tuple literals are likely unintentional and should be rewritten. This codemod rewrites the original `assert` statement by creating a new `assert` for each item in the original tuple.
+
+The changes from this codemod look like this:
+
+```diff
+- assert (1 == 1, 2 == 2)
++ assert 1 == 1
++ assert 2 == 2
+```

src/core_codemods/fix_assert_tuple.py

@@ -0,0 +1,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,
+                )
+            )

tests/codemods/base_codemod_test.py

@@ -18,8 +18,7 @@ class BaseCodemodTest:
     def setup_method(self):
         if isinstance(self.codemod, type):
             self.codemod = self.codemod()
-
-        self.file_context = None
+        self.changeset = []
 
     def run_and_assert(  # pylint: disable=too-many-arguments
         self,
@@ -52,6 +51,8 @@ def run_and_assert(  # pylint: disable=too-many-arguments
         self.codemod.apply(self.execution_context, files_to_check)
         changes = self.execution_context.get_results(self.codemod.id)
 
+        self.changeset = changes
+
         if input_code == expected:
             assert not changes
             return

tests/codemods/test_fix_assert_tuple.py

@@ -0,0 +1,89 @@
+import pytest
+from tests.codemods.base_codemod_test import BaseCodemodTest
+from core_codemods.fix_assert_tuple import FixAssertTuple
+
+
+class TestFixAssertTuple(BaseCodemodTest):
+    codemod = FixAssertTuple
+
+    def test_name(self):
+        assert self.codemod.name == "fix-assert-tuple"
+
+    @pytest.mark.parametrize(
+        "input_code,expected_output,change_count",
+        [
+            ("""assert (1,)""", """assert 1""", 1),
+            (
+                """assert ("one", Exception, [])""",
+                """\
+            assert "one"
+            assert Exception
+            assert []""",
+                3,
+            ),
+        ],
+    )
+    def test_change(self, tmpdir, input_code, expected_output, change_count):
+        self.run_and_assert(
+            tmpdir, input_code, expected_output, num_changes=change_count
+        )
+        for idx, change in enumerate(self.changeset[0].changes):
+            assert change.lineNumber == idx + 1
+
+    def test_change_line_pos(self, tmpdir):
+        input_code = """\
+        print(1)
+        assert (1, 2, )
+        print(2)
+        """
+        expected_output = """\
+        print(1)
+        assert 1
+        assert 2
+        print(2)
+        """
+
+        self.run_and_assert(tmpdir, input_code, expected_output, num_changes=2)
+        first_assert_line = 2
+        for idx, change in enumerate(self.changeset[0].changes):
+            assert change.lineNumber == idx + first_assert_line
+
+    def test_change_with_message(self, tmpdir):
+        input_code = """\
+        print(1)
+        assert (1, 2, ), "some message"
+        print(2)
+        """
+        expected_output = """\
+        print(1)
+        assert 1, "some message"
+        assert 2, "some message"
+        print(2)
+        """
+
+        self.run_and_assert(tmpdir, input_code, expected_output, num_changes=2)
+
+    @pytest.mark.parametrize(
+        "code",
+        [
+            "assert [1]",
+            "assert 1",
+            "assert ()",
+            "assert (1)",
+            "assert ('hi')",
+        ],
+    )
+    def test_no_change(self, tmpdir, code):
+        self.run_and_assert(tmpdir, code, code)
+
+    def test_exclude_line(self, tmpdir):
+        input_code = expected = """\
+        assert (1, 2)
+        """
+        lines_to_exclude = [1]
+        self.run_and_assert(
+            tmpdir,
+            input_code,
+            expected,
+            lines_to_exclude=lines_to_exclude,
+        )

tests/samples/fix_assert_tuple.py

@@ -0,0 +1 @@
+assert (1 == 1, 2 == 2)