Fix: Pycln does not skip some noqa imports (#88)

hadialqattan · Nov 12, 2021 · 621d025 · 621d025
1 parent ad23634
commit 621d025
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 15 deletions.
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
@@ -12,10 +12,11 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Added
 
-- [Add support for Python 3.10](https://github.com/hadialqattan/pycln/pull/81)
+- [Add support for Python 3.10 by @hadialqattan](https://github.com/hadialqattan/pycln/pull/81)
 
 ### Fixed
 
+- [Pycln does not skip imports that have "# nopycln: import" or "# noqa" on the last line by @hadialqattan](https://github.com/hadialqattan/pycln/pull/88)
 - [Pycln removes extra lines in import-from multiline case (shown bellow) by @hadialqattan](https://github.com/hadialqattan/pycln/pull/87)
 
   ```python3

diff --git a/pycln/utils/refactor.py b/pycln/utils/refactor.py
@@ -200,8 +200,11 @@ def _refactor(self, original_lines: List[str]) -> str:
             for node in type_:
 
                 # Skip any import that has `# noqa` or `# nopycln: import` comment.
-                lineno = node.location.start.line - 1
-                if regexu.skip_import(fixed_lines[lineno]):
+                s_lineno = node.location.start.line - 1
+                e_lineno = node.location.end.line - 1
+                if regexu.skip_import(fixed_lines[s_lineno]) or regexu.skip_import(
+                    fixed_lines[e_lineno]
+                ):
                     self.reporter.ignored_import(self._path, node)
                     continue
 

diff --git a/pycln/utils/scan.py b/pycln/utils/scan.py
@@ -184,7 +184,7 @@ def visit_Attribute(self, node: ast.Attribute):
         self._source_stats.attr_.add(node.attr)
 
     @recursive
-    def visit_MatchAs(self, node: 'ast.MatchAs'):
+    def visit_MatchAs(self, node: "ast.MatchAs"):  # type: ignore
         #: Support Match statement (PYTHON >= 3.10).
         #: PEP0634: https://www.python.org/dev/peps/pep-0634/
         self._source_stats.name_.add(node.name)

diff --git a/tests/test_refactor.py b/tests/test_refactor.py
@@ -263,17 +263,6 @@ def test_analyze(self, get_stats, get_stats_raise, expec_val):
             "_transform_return, expand_stars, mode, original_lines, expec_fixed_lines"
         ),
         [
-            pytest.param(
-                True,
-                None,
-                None,
-                ["import x # nopycln: import"],
-                False,
-                "not-matter",
-                ["import x, y # nopycln: import"],
-                ["import x, y # nopycln: import"],
-                id="nopycln",
-            ),
             pytest.param(
                 False,
                 (None, None),
@@ -400,6 +389,46 @@ def test_refactor(
                 fixed_code = self.session_maker._refactor(original_lines)
                 assert fixed_code == "".join(expec_fixed_lines)
 
+    @pytest.mark.parametrize(
+        ("endline_no, original_lines, expec_fixed_lines"),
+        [
+            pytest.param(
+                1,
+                ["import x # nopycln: import"],
+                ["import x # nopycln: import"],
+                id="nopycln - first line",
+            ),
+            pytest.param(
+                3,
+                ["import (\n", "    x\n", ") # nopycln: import"],
+                ["import (\n", "    x\n", ") # nopycln: import"],
+                id="nopycln - last line",
+            ),
+            pytest.param(
+                1,
+                ["import x # noqa"],
+                ["import x # noqa"],
+                id="noqa",
+            ),
+        ],
+    )
+    @mock.patch(MOCK % "Refactor._get_used_names")
+    def test_refactor_skipping(
+        self,
+        _get_used_names,
+        endline_no,
+        original_lines,
+        expec_fixed_lines,
+    ):
+        _get_used_names.return_value = set()
+        setattr(self.configs, "expand_stars", False)
+        node = Import(
+            NodeLocation((1, 0), endline_no), [ast.alias(name="x", asname=None)]
+        )
+        self.session_maker._import_stats = ImportStats({node}, set())
+        fixed_code = self.session_maker._refactor(original_lines)
+        assert fixed_code == "".join(expec_fixed_lines)
+
     @pytest.mark.parametrize(
         "_should_remove_return, node, is_star, expec_names",
         [