[pylint] Better diagnostic range (PLR1702)

crates/ruff_linter/resources/test/fixtures/pylint/too_many_nested_blocks_2.py

@@ -0,0 +1,33 @@
+def foo():
+    while a:                              # \
+        if b:                             # |
+            for c in range(3):            # | These should not be reported,
+                if d:                     # | as they don't exceed the max depth.
+                    while e:              # |
+                        if f:             # /
+
+                            for g in z:   # This statement is the first to exceed the limit.
+                                print(p)  # Thus, it is reported but not any of its substatements.
+                                pass      #
+
+                            with y:       # The former statement was already reported.
+                                print(x)  # Thus, reporting these is redundant.
+                            print(u)      #
+
+        else:                             # Other blocks of an ancestor statement
+            print(q)                      # are also not reported.
+
+
+def foo():
+    while a:                              # \
+        if b:                             # |
+            for c in range(3):            # | These should not be reported,
+                if d:                     # | as they don't exceed the max depth.
+                    while e:              # |
+                        if f:             # /
+
+                            if x == y:    # This statement is the first to exceed the limit.
+                                print(p)  # It is therefore reported.
+
+                            elif y > x:   # This block belongs to the same statement,
+                                print(p)  # and so it is not reported on its own.

crates/ruff_linter/src/rules/pylint/mod.rs

@@ -209,6 +209,7 @@ mod tests {
     #[test_case(Rule::SuperWithoutBrackets, Path::new("super_without_brackets.py"))]
     #[test_case(Rule::SelfOrClsAssignment, Path::new("self_or_cls_assignment.py"))]
     #[test_case(Rule::TooManyNestedBlocks, Path::new("too_many_nested_blocks.py"))]
+    #[test_case(Rule::TooManyNestedBlocks, Path::new("too_many_nested_blocks_2.py"))]
     #[test_case(Rule::DictIndexMissingItems, Path::new("dict_index_missing_items.py"))]
     #[test_case(Rule::DictIterMissingItems, Path::new("dict_iter_missing_items.py"))]
     #[test_case(

crates/ruff_linter/src/rules/pylint/rules/too_many_nested_blocks.rs

@@ -1,8 +1,12 @@
+use std::ptr;
+
 use ast::ExceptHandler;
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::{self as ast, Stmt};
-use ruff_text_size::Ranged;
+use ruff_python_semantic::SemanticModel;
+use ruff_source_file::LineRanges;
+use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
 
@@ -37,8 +41,10 @@ impl Violation for TooManyNestedBlocks {
 
 /// PLR1702
 pub(crate) fn too_many_nested_blocks(checker: &mut Checker, stmt: &Stmt) {
+    let semantic = checker.semantic();
+
     // Only enforce nesting within functions or methods.
-    if !checker.semantic().current_scope().kind.is_function() {
+    if !semantic.current_scope().kind.is_function() {
         return;
     }
 
@@ -48,38 +54,49 @@ pub(crate) fn too_many_nested_blocks(checker: &mut Checker, stmt: &Stmt) {
         return;
     }
 
+    if !current_stmt_is_first_within_parent(semantic) {
+        return;
+    }
+
     let max_nested_blocks = checker.settings.pylint.max_nested_blocks;
 
+    let Some(current_id) = semantic.current_statement_id() else {
+        return;
+    };
+
     // Traverse up the hierarchy, identifying the root node and counting the number of nested
     // blocks between the root and this leaf.
-    let (count, root_id) =
-        checker
-            .semantic()
+    let (count, encountered_root) =
+        semantic
             .current_statement_ids()
-            .fold((0, None), |(count, ancestor_id), id| {
-                let stmt = checker.semantic().statement(id);
+            .fold((0, false), |(count, encountered_root), id| {
+                let stmt = semantic.statement(id);
                 if is_nested_block(stmt) {
-                    (count + 1, Some(id))
+                    (count + 1, true)
                 } else {
-                    (count, ancestor_id)
+                    (count, encountered_root)
                 }
             });
 
-    let Some(root_id) = root_id else {
+    if !encountered_root {
         return;
-    };
+    }
 
     // If the number of nested blocks is less than the maximum, we don't want to emit a diagnostic.
     if count <= max_nested_blocks {
         return;
     }
 
+    let current_stmt_start = semantic.statement(current_id).start();
+    let current_stmt_line_start = checker.locator().line_start(current_stmt_start);
+    let indentation_range = TextRange::new(current_stmt_line_start, current_stmt_start);
+
     checker.diagnostics.push(Diagnostic::new(
         TooManyNestedBlocks {
             nested_blocks: count,
             max_nested_blocks,
         },
-        checker.semantic().statement(root_id).range(),
+        indentation_range,
     ));
 }
 
@@ -91,7 +108,7 @@ fn is_nested_block(stmt: &Stmt) -> bool {
     )
 }
 
-/// Returns `true` if the given statement is a leaf node.
+/// Returns `true` if the given statement is not a leaf node.
 fn has_nested_block(stmt: &Stmt) -> bool {
     match stmt {
         Stmt::If(ast::StmtIf {
@@ -130,3 +147,25 @@ fn has_nested_block(stmt: &Stmt) -> bool {
         _ => false,
     }
 }
+
+fn current_stmt_is_first_within_parent(semantic: &SemanticModel) -> bool {
+    let Some(parent) = semantic.current_statement_parent() else {
+        return false;
+    };
+    let current = semantic.current_statement();
+
+    match parent {
+        Stmt::If(ast::StmtIf { body, .. })
+        | Stmt::While(ast::StmtWhile { body, .. })
+        | Stmt::For(ast::StmtFor { body, .. })
+        | Stmt::Try(ast::StmtTry { body, .. })
+        | Stmt::With(ast::StmtWith { body, .. }) => {
+            let [first, ..] = &body[..] else {
+                return false;
+            };
+
+            ptr::eq(first, current)
+        }
+        _ => false,
+    }
+}

crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR1702_too_many_nested_blocks.py.snap

@@ -1,17 +1,12 @@
 ---
 source: crates/ruff_linter/src/rules/pylint/mod.rs
 ---
-too_many_nested_blocks.py:2:5: PLR1702 Too many nested blocks (6 > 5)
+too_many_nested_blocks.py:8:1: PLR1702 Too many nested blocks (6 > 5)
    |
- 1 |   def correct_fruits(fruits) -> bool:
- 2 | /     if len(fruits) > 1:  # PLR1702
- 3 | |         if "apple" in fruits:
- 4 | |             if "orange" in fruits:
- 5 | |                 count = fruits["orange"]
- 6 | |                 if count % 2:
- 7 | |                     if "kiwi" in fruits:
- 8 | |                         if count == 2:
- 9 | |                             return True
-   | |_______________________________________^ PLR1702
-10 |       return False
+ 6 |                 if count % 2:
+ 7 |                     if "kiwi" in fruits:
+ 8 |                         if count == 2:
+   | ^^^^^^^^^^^^^^^^^^^^^^^^ PLR1702
+ 9 |                             return True
+10 |     return False
    |

crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR1702_too_many_nested_blocks_2.py.snap

@@ -0,0 +1,21 @@
+---
+source: crates/ruff_linter/src/rules/pylint/mod.rs
+---
+too_many_nested_blocks_2.py:9:1: PLR1702 Too many nested blocks (7 > 5)
+   |
+ 7 |                         if f:             # /
+ 8 |
+ 9 |                             for g in z:   # This statement is the first to exceed the limit.
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PLR1702
+10 |                                 print(p)  # Thus, it is reported but not any of its substatements.
+11 |                                 pass      #
+   |
+
+too_many_nested_blocks_2.py:29:1: PLR1702 Too many nested blocks (7 > 5)
+   |
+27 |                         if f:             # /
+28 |
+29 |                             if x == y:    # This statement is the first to exceed the limit.
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PLR1702
+30 |                                 print(p)  # It is therefore reported.
+   |