Added support for disable unused import pylint pragma

pyproject.toml

@@ -9,11 +9,12 @@ requires-python = ">=3.9.0"
 readme = "README.md"
 license = {file = "LICENSE"}
 dependencies = [
-    "semgrep~=1.41.0",
-    "PyYAML~=6.0.0",
-    "libcst~=1.0.0",
-    "isort~=5.12.0",
     "dependency-manager @ git+https://github.com/pixee/python-dependency-manager#egg=dependency-manager",
+    "isort~=5.12.0",
+    "libcst~=1.0.0",
+    "pylint~=2.17.0",
+    "PyYAML~=6.0.0",
+    "semgrep~=1.41.0",
 ]
 
 [project.scripts]

requirements/lint.txt

@@ -1,4 +1,3 @@
 black==23.9.*
 mypy==1.5.*
-pylint==2.17.*
 -r test.txt

src/codemodder/codemods/remove_unused_imports.py

@@ -18,6 +18,9 @@
 import libcst as cst
 from libcst.codemod import Codemod, CodemodContext
 import re
+from pylint.utils.pragma_parser import parse_pragma
+
+NOQA_PATTERN = re.compile(r"^#\s*noqa")
 
 
 class RemoveUnusedImports(BaseCodemod, Codemod):
@@ -48,7 +51,7 @@ def transform_module_impl(self, tree: cst.Module) -> cst.Module:
         for import_alias, importt in gather_unused_visitor.unused_imports:
             pos = self.get_metadata(PositionProvider, import_alias)
             if self.filter_by_path_includes_or_excludes(pos):
-                if not self._has_noqa_comment(importt):
+                if not self._is_disabled_by_linter(importt):
                     self.file_context.codemod_changes.append(
                         Change(pos.start.line, self.CHANGE_DESCRIPTION).to_json()
                     )
@@ -66,22 +69,25 @@ def filter_by_path_includes_or_excludes(self, pos_to_match):
             return any(match_line(pos_to_match, line) for line in self.line_include)
         return True
 
-    def _has_noqa_comment(self, node):
+    def _is_disabled_by_linter(self, node):
         """
-        Check if the import has a #noqa comment attached to it
+        Check if the import has a #noqa or # pylint: disable(-next)=unused_imports comment attached to it.
         """
         parent = self.get_metadata(ParentNodeProvider, node)
         if parent and matchers.matches(parent, matchers.SimpleStatementLine()):
             stmt = ensure_type(parent, cst.SimpleStatementLine)
-            pattern = re.compile(r"^#\s*noqa")
 
             # has a trailing comment string
             trailing_comment_string = (
                 stmt.trailing_whitespace.comment.value
                 if stmt.trailing_whitespace.comment
                 else None
             )
-            if trailing_comment_string and pattern.match(trailing_comment_string):
+            if trailing_comment_string and NOQA_PATTERN.match(trailing_comment_string):
+                return True
+            if trailing_comment_string and _is_pylint_disable_unused_imports(
+                trailing_comment_string
+            ):
                 return True
 
             # has a comment right above it
@@ -94,10 +100,37 @@ def _has_noqa_comment(self, node):
                     ]
                 ),
             ):
-                if pattern.match(stmt.leading_lines[-1].comment.value):
+                comment_string = stmt.leading_lines[-1].comment.value
+                if NOQA_PATTERN.match(comment_string):
+                    return True
+                if comment_string and _is_pylint_disable_next_unused_imports(
+                    comment_string
+                ):
                     return True
         return False
 
 
 def match_line(pos, line):
     return pos.start.line == line and pos.end.line == line
+
+
+def _is_pylint_disable_unused_imports(comment: str) -> bool:
+    parsed = parse_pragma(comment)
+    for p in parsed:
+        if p.action == "disable" and (
+            "unused-import" in p.messages or "W0611" in p.messages
+        ):
+            return True
+    return False
+
+
+def _is_pylint_disable_next_unused_imports(comment: str) -> bool:
+    parsed = parse_pragma(comment)
+    for p in parsed:
+        print(p.action)
+        print(p.messages)
+        if p.action == "disable-next" and (
+            "unused-import" in p.messages or "W0611" in p.messages
+        ):
+            return True
+    return False

tests/codemods/test_remove_unused_imports.py

@@ -89,3 +89,15 @@ def test_dont_remove_if_noqa_trailing(self, tmpdir):
         before = "import a\nimport b # noqa\na()"
         self.run_and_assert(tmpdir, before, before)
         assert len(self.file_context.codemod_changes) == 0
+
+    def test_dont_remove_if_pylint_disable(self, tmpdir):
+        before = "import a\nimport b # pylint: disable=W0611\na()"
+        self.run_and_assert(tmpdir, before, before)
+        assert len(self.file_context.codemod_changes) == 0
+
+    def test_dont_remove_if_pylint_disable_next(self, tmpdir):
+        before = (
+            "import a\n#   pylint: disable-next=no-member, unused-import\nimport b\na()"
+        )
+        self.run_and_assert(tmpdir, before, before)
+        assert len(self.file_context.codemod_changes) == 0