[pylint] Better diagnostic range (PLR1702)

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

@@ -0,0 +1,59 @@
+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 nested statement is the first to have
+                                print(p)          # a substatement that exceed the limit.
+                                pass              # It is reported but not any of its substatements.
+
+                            with y:               # The former statement was already reported.
+                                print(x)          # Thus, reporting these is redundant.
+                            print(u)              #
+
+        else:                                     # \
+            print(q)                              # |
+            if h:                                 # | Other blocks of an ancestor statement
+                if i:                             # | are also not reported.
+                    if j:                         # |
+                        if k:                     # /
+
+                            if l:                 # Unless they themselves have a branch
+                                print()           # that exceeds the limit.
+
+
+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.
+
+
+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:                      # |
+                        print()                   # |
+                        if f:                     # /
+
+                            if g:                 # This statement is the first nested statement to exceed the limit.
+                                print()           # However, it is not reported due to a nested block below.
+
+                            else:                 # This part of the block does not participate
+                                print()           # in the tree in question and is thus not reported.
+
+                                if i:             # This branch does.
+                                    print()       #

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,10 @@
-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_ast::{self as ast, ExceptHandler, Stmt};
+use ruff_python_semantic::SemanticModel;
+use ruff_source_file::LineRanges;
+use ruff_text_size::{Ranged, TextRange};
+use std::{iter, ptr};
 
 use crate::checkers::ast::Checker;
 
@@ -37,8 +39,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 +52,49 @@ pub(crate) fn too_many_nested_blocks(checker: &mut Checker, stmt: &Stmt) {
         return;
     }
 
+    if !current_stmt_is_first_nested_within_block(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 +106,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 +145,51 @@ fn has_nested_block(stmt: &Stmt) -> bool {
         _ => false,
     }
 }
+
+fn current_stmt_is_first_nested_within_block(semantic: &SemanticModel) -> bool {
+    let Some(parent) = semantic.current_statement_parent() else {
+        return false;
+    };
+    let current = semantic.current_statement();
+
+    let current_is_first_nested = |block: &Vec<Stmt>| {
+        let Some(first_nested) = block.iter().find(|stmt| is_nested_block(stmt)) else {
+            return false;
+        };
+
+        ptr::eq(first_nested, current)
+    };
+
+    match parent {
+        Stmt::With(ast::StmtWith { body, .. }) => current_is_first_nested(body),
+
+        Stmt::If(ast::StmtIf {
+            body,
+            elif_else_clauses,
+            ..
+        }) => iter::once(body)
+            .chain(elif_else_clauses.iter().map(|clause| &clause.body))
+            .any(current_is_first_nested),
+
+        Stmt::While(ast::StmtWhile { body, orelse, .. })
+        | Stmt::For(ast::StmtFor { body, orelse, .. }) => iter::once(body)
+            .chain(iter::once(orelse))
+            .any(current_is_first_nested),
+
+        Stmt::Try(ast::StmtTry {
+            body,
+            handlers,
+            orelse,
+            finalbody,
+            ..
+        }) => handlers
+            .iter()
+            .map(|ExceptHandler::ExceptHandler(handler)| &handler.body)
+            .chain(iter::once(body))
+            .chain(iter::once(orelse))
+            .chain(iter::once(finalbody))
+            .any(current_is_first_nested),
+
+        _ => 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,39 @@
+---
+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 nested statement is the first to have
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PLR1702
+10 |                                 print(p)          # a substatement that exceed the limit.
+11 |                                 pass              # It is reported but not any of its substatements.
+   |
+
+too_many_nested_blocks_2.py:24:1: PLR1702 Too many nested blocks (7 > 5)
+   |
+22 |                         if k:                     # /
+23 |
+24 |                             if l:                 # Unless they themselves have a branch
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PLR1702
+25 |                                 print()           # that exceeds the limit.
+   |
+
+too_many_nested_blocks_2.py:36:1: PLR1702 Too many nested blocks (7 > 5)
+   |
+34 |                         if f:                     # /
+35 |
+36 |                             if x == y:            # This statement is the first to exceed the limit.
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PLR1702
+37 |                                 print(p)          # It is therefore reported.
+   |
+
+too_many_nested_blocks_2.py:58:1: PLR1702 Too many nested blocks (8 > 5)
+   |
+56 |                                 print()           # in the tree in question and is thus not reported.
+57 |
+58 |                                 if i:             # This branch does.
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PLR1702
+59 |                                     print()       #
+   |