match by codemod name or id

pixee · Sep 7, 2023 · 9bd6078 · 9bd6078
1 parent 1556b10
commit 9bd6078
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/src/codemodder/codemods/__init__.py b/src/codemodder/codemods/__init__.py
@@ -53,10 +53,12 @@ def match_codemods(codemod_include: list, codemod_exclude: list) -> dict:
             name: codemod
             for codemod in DEFAULT_CODEMODS
             if (name := codemod.name()) not in codemod_exclude
+            and (_ := codemod.id()) not in codemod_exclude
         }
 
     return {
         name: codemod
         for codemod in DEFAULT_CODEMODS
         if (name := codemod.name()) in codemod_include
+        or (_ := codemod.id()) in codemod_include
     }
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -110,6 +110,20 @@ def test_no_codemods_to_run(self, mock_semgrep_run):
         assert exit_code == 0
         mock_semgrep_run.assert_not_called()
 
+    @pytest.mark.parametrize("codemod", ["secure-random", "pixee:python/secure-random"])
+    @mock.patch("codemodder.__main__.compile_results")
+    def test_run_codemod_name_or_id(self, mock_compile_results, codemod):
+        args = [
+            "tests/samples/",
+            "--output",
+            "here.txt",
+            f"--codemod-include={codemod}",
+        ]
+
+        exit_code = run(parse_args(args), args)
+        assert exit_code == 0
+        mock_compile_results.assert_called()
+
     @mock.patch("codemodder.semgrep.results_by_path_and_rule_id")
     @mock.patch("codemodder.__main__.semgrep_run", side_effect=semgrep_run)
     @mock.patch("codemodder.semgrep.subprocess.run")