A bit of cleanup

pixee · Feb 13, 2024 · b3cf5cb · b3cf5cb
1 parent 453aa9c
commit b3cf5cb
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/src/codemodder/codemods/check_annotations.py b/src/codemodder/codemods/check_annotations.py
@@ -67,7 +67,7 @@ def _process_simple_statement_line(self, stmt: cst.SimpleStatementLine) -> bool:
 
     def is_disabled_by_linter(self, node: cst.CSTNode) -> bool:
         """
-        Check if the import has a #noqa or # pylint: disable(-next)=unused_imports comment attached to it.
+        Check if the import has a #noqa or # pylint: disable(-next) comment attached to it.
         """
         match self.get_metadata(ParentNodeProvider, node):
             case cst.SimpleStatementLine() as stmt:
@@ -120,7 +120,7 @@ def is_disabled_by_annotations(
     messages: list[str],
 ) -> bool:
     """
-    Check if the import has a #noqa or # pylint: disable(-next)=unused_imports comment attached to it.
+    Check if the import has a #noqa or # pylint: disable(-next) comment attached to it.
     """
     visitor = _GatherCommentNodes(metadata, messages)
     node.visit(visitor)

diff --git a/src/core_codemods/remove_unused_imports.py b/src/core_codemods/remove_unused_imports.py
@@ -30,6 +30,8 @@ class RemoveUnusedImports(SimpleCodemod):
         ParentNodeProvider,
     )
 
+    IGNORE_ANNOTATIONS = ["unused-import", "F401", "W0611"]
+
     def transform_module_impl(self, tree: cst.Module) -> cst.Module:
         # Do nothing in __init__.py files
         if self.file_context.file_path.name == "__init__.py":
@@ -42,7 +44,9 @@ def transform_module_impl(self, tree: cst.Module) -> cst.Module:
             pos = self.get_metadata(PositionProvider, import_alias)
             if self.filter_by_path_includes_or_excludes(pos):
                 if not is_disabled_by_annotations(
-                    importt, self.metadata, messages=["unused-import", "W0611", "F401"]  # type: ignore
+                    importt,
+                    self.metadata,  # type: ignore
+                    messages=self.IGNORE_ANNOTATIONS,
                 ):
                     self.file_context.codemod_changes.append(
                         Change(pos.start.line, self.change_description)

diff --git a/src/core_codemods/subprocess_shell_false.py b/src/core_codemods/subprocess_shell_false.py
@@ -35,6 +35,7 @@ class SubprocessShellFalse(SimpleCodemod, NameResolutionMixin):
     ]
 
     METADATA_DEPENDENCIES = SimpleCodemod.METADATA_DEPENDENCIES + (ParentNodeProvider,)
+    IGNORE_ANNOTATIONS = ["S603"]
 
     def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
         if not self.filter_by_path_includes_or_excludes(
@@ -52,7 +53,7 @@ def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
                 ) and not is_disabled_by_annotations(
                     original_node,
                     self.metadata,  # type: ignore
-                    messages=["S603"],
+                    messages=self.IGNORE_ANNOTATIONS,
                 ):
                     self.report_change(original_node)
                     new_args = self.replace_args(