From 328b24ed35d36217a9d5cb1cff5c80d74e085ace Mon Sep 17 00:00:00 2001
From: andrecs <12188364+andrecsilva@users.noreply.github.com>
Date: Mon, 25 Sep 2023 09:07:53 -0300
Subject: [PATCH] Added support for disable unused import pylint pragma

---
 pyproject.toml                                |  9 ++--
 requirements/lint.txt                         |  1 -
 .../codemods/remove_unused_imports.py         | 45 ++++++++++++++++---
 tests/codemods/test_remove_unused_imports.py  | 12 +++++
 4 files changed, 56 insertions(+), 11 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ae906c6d..a01a6125 100644
--- a/pyproject.toml
+++ b/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]
diff --git a/requirements/lint.txt b/requirements/lint.txt
index e3e1ed51..ffaa081b 100644
--- a/requirements/lint.txt
+++ b/requirements/lint.txt
@@ -1,4 +1,3 @@
 black==23.9.*
 mypy==1.5.*
-pylint==2.17.*
 -r test.txt
diff --git a/src/codemodder/codemods/remove_unused_imports.py b/src/codemodder/codemods/remove_unused_imports.py
index 880b0df1..9ee5e2fd 100644
--- a/src/codemodder/codemods/remove_unused_imports.py
+++ b/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,14 +69,13 @@ 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 = (
@@ -81,7 +83,11 @@ def _has_noqa_comment(self, node):
                 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
diff --git a/tests/codemods/test_remove_unused_imports.py b/tests/codemods/test_remove_unused_imports.py
index d2c74311..ce8d9fec 100644
--- a/tests/codemods/test_remove_unused_imports.py
+++ b/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