Support block-level wrong-import-position pragma suppression

Signed-off-by: Emmanuel Ferdman <[email protected]>

Author: emmanuel-ferdman

doc/whatsnew/fragments/10589.bugfix

@@ -1,3 +1,3 @@
-Scope ``wrong-import-position`` pragma to specific line only.
+Allow ``wrong-import-position`` pragma on non-import lines to suppress following imports until the next non-import statement.
 
 Closes #10589

pylint/checkers/imports.py

@@ -447,6 +447,7 @@ def __init__(self, linter: PyLinter) -> None:
         self.import_graph: defaultdict[str, set[str]] = defaultdict(set)
         self._imports_stack: list[tuple[ImportNode, str]] = []
         self._first_non_import_node = None
+        self._non_import_nodes: list = []
         self._module_pkg: dict[Any, Any] = (
             {}
         )  # mapping of modules to the pkg they belong in
@@ -608,6 +609,7 @@ def leave_module(self, node: nodes.Module) -> None:
 
         self._imports_stack = []
         self._first_non_import_node = None
+        self._non_import_nodes = []
 
     def compute_first_non_import_node(
         self,
@@ -621,12 +623,7 @@ def compute_first_non_import_node(
             | nodes.Try
         ),
     ) -> None:
-        # if the node does not contain an import instruction, and if it is the
-        # first node of the module, keep a track of it (all the import positions
-        # of the module will be compared to the position of this first
-        # instruction)
-        if self._first_non_import_node:
-            return
+        # Track non-import nodes at module level to check import positions
         if not isinstance(node.parent, nodes.Module):
             return
         if isinstance(node, nodes.Try) and any(
@@ -644,7 +641,11 @@ def compute_first_non_import_node(
             ]
             if all(valid_targets):
                 return
-        self._first_non_import_node = node
+
+        if not self._first_non_import_node:
+            self._first_non_import_node = node
+
+        self._non_import_nodes.append(node)
 
     visit_try = visit_assignattr = visit_assign = visit_ifexp = visit_comprehension = (
         visit_expr
@@ -653,12 +654,7 @@ def compute_first_non_import_node(
     def visit_functiondef(
         self, node: nodes.FunctionDef | nodes.While | nodes.For | nodes.ClassDef
     ) -> None:
-        # If it is the first non import instruction of the module, record it.
-        if self._first_non_import_node:
-            return
-
-        # Check if the node belongs to an `If` or a `Try` block. If they
-        # contain imports, skip recording this node.
+        # Record non-import instruction unless inside an If/Try block that contains imports
         if not isinstance(node.parent.scope(), nodes.Module):
             return
 
@@ -670,7 +666,10 @@ def visit_functiondef(
             if any(root.nodes_of_class((nodes.Import, nodes.ImportFrom))):
                 return
 
-        self._first_non_import_node = node
+        if not self._first_non_import_node:
+            self._first_non_import_node = node
+
+        self._non_import_nodes.append(node)
 
     visit_classdef = visit_for = visit_while = visit_functiondef
 
@@ -699,17 +698,37 @@ def _check_position(self, node: ImportNode) -> None:
 
         Send a message  if `node` comes before another instruction
         """
-        # if a first non-import instruction has already been encountered,
-        # it means the import comes after it and therefore is not well placed
+        # Check if import comes after a non-import statement
         if self._first_non_import_node:
-            if self.linter.is_message_enabled("wrong-import-position", node.fromlineno):
-                self.add_message(
-                    "wrong-import-position", node=node, args=node.as_string()
-                )
-            else:
+            # Check for inline pragma on the import line
+            if not self.linter.is_message_enabled("wrong-import-position", node.fromlineno):
                 self.linter.add_ignored_message(
                     "wrong-import-position", node.fromlineno, node
                 )
+                return
+
+            # Check for pragma on the preceding non-import statement
+            most_recent_non_import = None
+            for non_import_node in self._non_import_nodes:
+                if non_import_node.fromlineno < node.fromlineno:
+                    most_recent_non_import = non_import_node
+                else:
+                    break
+
+            if most_recent_non_import:
+                check_line = most_recent_non_import.fromlineno
+                if not self.linter.is_message_enabled("wrong-import-position", check_line):
+                    self.linter.add_ignored_message(
+                        "wrong-import-position", check_line, most_recent_non_import
+                    )
+                    self.linter.add_ignored_message(
+                        "wrong-import-position", node.fromlineno, node
+                    )
+                    return
+
+            self.add_message(
+                "wrong-import-position", node=node, args=node.as_string()
+            )
 
     def _record_import(
         self,

tests/functional/d/disable_wrong_import_position.py

@@ -1,10 +1,11 @@
-"""Checks that disabling 'wrong-import-position' only affects the specific line.
-
-A pragma on a non-import statement should not affect subsequent import statements.
-This demonstrates the correct behavior after fixing the bug.
-"""
+"""Test wrong-import-position pragma on non-import statement."""
 # pylint: disable=unused-import
 
-CONSTANT = True  # pylint: disable=wrong-import-position
+import os
+import sys
+
+CONSTANT_A = False  # pylint: disable=wrong-import-position
+import time
 
-import sys  # [wrong-import-position]
+CONSTANT_B = True
+import logging  # [wrong-import-position]

tests/functional/d/disable_wrong_import_position.txt

@@ -1 +1 @@
-wrong-import-position:10:0:10:10::"Import ""import sys"" should be placed at the top of the module":UNDEFINED
+wrong-import-position:11:0:11:14::"Import ""import logging"" should be placed at the top of the module":UNDEFINED

tests/functional/w/wrong_import_position_pragma_scope.py

@@ -1,23 +1,16 @@
-"""Test that wrong-import-position pragma suppression is correctly scoped."""
-# pylint: disable=unused-import,invalid-name
+"""Test wrong-import-position pragma scoping."""
+# pylint: disable=unused-import
 
-import logging
+import os
 import sys
 
-# This pragma should not affect subsequent import statements
-logger = logging.getLogger()  # pylint: disable=wrong-import-position
-logging.basicConfig(level='DEBUG')
+# Pragma on non-import suppresses following imports until next non-import
+CONSTANT_A = False  # pylint: disable=wrong-import-position
+import time
 
-logger.debug('importing modules...')
-import os  # [wrong-import-position]
-import pathlib  # [wrong-import-position]
-import random  # [wrong-import-position]
-logger.debug('done importing')
+CONSTANT_B = True
+import logging  # [wrong-import-position]
 
-# Test that pragma on import line works correctly (this import should not be flagged)
-constant_var = "test"
+# Inline pragma on import line
+CONSTANT_C = 42
 import json  # pylint: disable=wrong-import-position
-
-# Test that subsequent imports are not affected by the pragma above
-import csv  # [wrong-import-position]
-import re  # [wrong-import-position]

tests/functional/w/wrong_import_position_pragma_scope.txt

@@ -1,5 +1 @@
-wrong-import-position:12:0:12:9::"Import ""import os"" should be placed at the top of the module":UNDEFINED
-wrong-import-position:13:0:13:14::"Import ""import pathlib"" should be placed at the top of the module":UNDEFINED
-wrong-import-position:14:0:14:13::"Import ""import random"" should be placed at the top of the module":UNDEFINED
-wrong-import-position:22:0:22:10::"Import ""import csv"" should be placed at the top of the module":UNDEFINED
-wrong-import-position:23:0:23:9::"Import ""import re"" should be placed at the top of the module":UNDEFINED
+wrong-import-position:12:0:12:14::"Import ""import logging"" should be placed at the top of the module":UNDEFINED