new codemod to combine startswith/endswith

integration_tests/test_combine_startswith_endswith.py

@@ -0,0 +1,16 @@
+from core_codemods.combine_startswith_endswith import CombineStartswithEndswith
+from integration_tests.base_test import (
+    BaseIntegrationTest,
+    original_and_expected_from_code_path,
+)
+
+
+class TestCombineStartswithEndswith(BaseIntegrationTest):
+    codemod = CombineStartswithEndswith
+    code_path = "tests/samples/combine_startswith_endswith.py"
+    original_code, expected_new_code = original_and_expected_from_code_path(
+        code_path, [(1, 'if x.startswith(("foo", "bar")):\n')]
+    )
+    expected_diff = '--- \n+++ \n@@ -1,3 +1,3 @@\n x = \'foo\'\n-if x.startswith("foo") or x.startswith("bar"):\n+if x.startswith(("foo", "bar")):\n     print("Yes")\n'
+    expected_line_change = "2"
+    change_description = CombineStartswithEndswith.CHANGE_DESCRIPTION

src/codemodder/scripts/generate_docs.py

@@ -178,6 +178,10 @@ class DocMetadata:
         importance="Low",
         guidance_explained="Since literals and new objects have their own identities, comparisons against them using `is` operators are most likely a bug and thus we deem the change safe.",
     ),
+    "combine-startswith-endswith": DocMetadata(
+        importance="Low",
+        guidance_explained="Combining two `startswith` or `endswith` calls is safe and cleans up code.",
+    ),
 }
 
 

src/core_codemods/__init__.py

@@ -38,6 +38,7 @@
 from .sql_parameterization import SQLQueryParameterization
 from .exception_without_raise import ExceptionWithoutRaise
 from .literal_or_new_object_identity import LiteralOrNewObjectIdentity
+from .combine_startswith_endswith import CombineStartswithEndswith
 
 registry = CodemodCollection(
     origin="pixee",
@@ -82,5 +83,6 @@
         FlaskJsonResponseType,
         ExceptionWithoutRaise,
         LiteralOrNewObjectIdentity,
+        CombineStartswithEndswith,
     ],
 )

src/core_codemods/combine_startswith_endswith.py

@@ -0,0 +1,49 @@
+import libcst as cst
+from libcst import matchers as m
+from codemodder.codemods.api import BaseCodemod, ReviewGuidance
+from codemodder.codemods.utils_mixin import NameResolutionMixin
+
+
+class CombineStartswithEndswith(BaseCodemod, NameResolutionMixin):
+    NAME = "combine-startswith-endswith"
+    SUMMARY = "Combine Two `startswith` or `endswith` Calls"
+    REVIEW_GUIDANCE = ReviewGuidance.MERGE_WITHOUT_REVIEW
+    DESCRIPTION = "Combine two calls to either `startswith` or `endswith` joined by an `or` statement."
+    REFERENCES: list = []
+
+    def leave_BooleanOperation(
+        self, original_node: cst.BooleanOperation, updated_node: cst.BooleanOperation
+    ) -> cst.CSTNode:
+        if self.matches_startswith_endswith_or_pattern(original_node):
+            left_call = cst.ensure_type(updated_node.left, cst.Call)
+            right_call = cst.ensure_type(updated_node.right, cst.Call)
+
+            self.report_change(original_node)
+
+            new_arg = cst.Arg(
+                value=cst.Tuple(
+                    elements=[
+                        cst.Element(value=left_call.args[0].value),
+                        cst.Element(value=right_call.args[0].value),
+                    ]
+                )
+            )
+
+            return cst.Call(func=left_call.func, args=[new_arg])
+
+        return updated_node
+
+    def matches_startswith_endswith_or_pattern(
+        self, node: cst.BooleanOperation
+    ) -> bool:
+        # Match the pattern: x.startswith("...") or x.startswith("...")
+        # and the same but with endswith
+        call = m.Call(
+            func=m.Attribute(
+                value=m.Name(), attr=m.Name("startswith") | m.Name("endswith")
+            ),
+            args=[m.Arg(value=m.SimpleString())],
+        )
+        return m.matches(
+            node, m.BooleanOperation(left=call, operator=m.Or(), right=call)
+        )

src/core_codemods/docs/pixee_python_combine-startswith-endswith.md

@@ -0,0 +1,11 @@
+This codemod updates places where two separate calls to either `startswith` or `endswith` are joined by an `or` statement. 
+The replacement code takes advantage of the fact that these two methods can accept a tuple as an argument.
+
+The changes from this codemod look like this:
+
+```diff
+  x = 'foo'
+- if x.startswith("foo") or x.startswith("bar"):
++ if x.startswith(("foo", "bar")):
+     ...
+```

tests/codemods/test_combine_startswith_endswith.py

@@ -0,0 +1,39 @@
+import pytest
+from tests.codemods.base_codemod_test import BaseCodemodTest
+from core_codemods.combine_startswith_endswith import CombineStartswithEndswith
+from textwrap import dedent
+
+each_func = pytest.mark.parametrize("func", ["startswith", "endswith"])
+
+
+class TestCombineStartswithEndswith(BaseCodemodTest):
+    codemod = CombineStartswithEndswith
+
+    def test_name(self):
+        assert self.codemod.name() == "combine-startswith-endswith"
+
+    @each_func
+    def test_combine(self, tmpdir, func):
+        input_code = f"""\
+        x = "foo"
+        x.{func}("foo") or x.{func}("f")
+        """
+        expected = f"""\
+        x = "foo"
+        x.{func}(("foo", "f"))
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    @pytest.mark.parametrize(
+        "code",
+        [
+            "x.startswith('foo')",
+            "x.startswith(('f', 'foo'))",
+            "x.startswith('foo') and x.startswith('f')",
+            "x.startswith('foo') and x.startswith('f') or True",
+        ],
+    )
+    def test_no_change(self, tmpdir, code):
+        self.run_and_assert(tmpdir, code, code)
+        assert len(self.file_context.codemod_changes) == 0

tests/samples/combine_startswith_endswith.py

@@ -0,0 +1,3 @@
+x = 'foo'
+if x.startswith("foo") or x.startswith("bar"):
+    print("Yes")