diff --git a/crates/ruff_linter/resources/test/fixtures/pylint/too_many_branches.py b/crates/ruff_linter/resources/test/fixtures/pylint/too_many_branches.py index b19314dc972ec..2bc360ca13bef 100644 --- a/crates/ruff_linter/resources/test/fixtures/pylint/too_many_branches.py +++ b/crates/ruff_linter/resources/test/fixtures/pylint/too_many_branches.py @@ -2,6 +2,8 @@ Test for too many branches. Taken from the pylint source 2023-02-03 """ +from contextlib import suppress + # pylint: disable=using-constant-test def wrong(): # [too-many-branches] """ Has too many branches. """ @@ -70,3 +72,33 @@ def nested_1(): pass elif 7: pass + +def with_statement_wrong(): + """statements inside the with statement should get counted""" + with suppress(Exception): + if 1: + pass + elif 1: + pass + elif 1: + pass + elif 1: + pass + elif 1: + pass + elif 1: + pass + try: + pass + finally: + pass + if 2: + pass + while True: + pass + if 1: + pass + elif 2: + pass + elif 3: + pass diff --git a/crates/ruff_linter/src/rules/pylint/rules/too_many_branches.rs b/crates/ruff_linter/src/rules/pylint/rules/too_many_branches.rs index ca14e12dcc5e2..9a1ea12053672 100644 --- a/crates/ruff_linter/src/rules/pylint/rules/too_many_branches.rs +++ b/crates/ruff_linter/src/rules/pylint/rules/too_many_branches.rs @@ -107,6 +107,8 @@ fn num_branches(stmts: &[Stmt]) -> usize { .map(|case| num_branches(&case.body)) .sum::() } + Stmt::With(ast::StmtWith { body, .. }) => num_branches(body), // The `with` statement + // is not considered a branch but the statements inside the `with` should be counted Stmt::For(ast::StmtFor { body, orelse, .. }) | Stmt::While(ast::StmtWhile { body, orelse, .. }) => { 1 + num_branches(body) @@ -265,4 +267,18 @@ finally: test_helper(source, 5)?; Ok(()) } + + #[test] + fn with_statement() -> Result<()> { + let source: &str = r" +with suppress(Exception): + if x == 0: # 2 + return + else: + return +"; + + test_helper(source, 2)?; + Ok(()) + } } diff --git a/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR0912_too_many_branches.py.snap b/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR0912_too_many_branches.py.snap index 690939ac3ff5e..559e5d6210760 100644 --- a/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR0912_too_many_branches.py.snap +++ b/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR0912_too_many_branches.py.snap @@ -1,14 +1,21 @@ --- source: crates/ruff_linter/src/rules/pylint/mod.rs --- -too_many_branches.py:6:5: PLR0912 Too many branches (13 > 12) - | -4 | """ -5 | # pylint: disable=using-constant-test -6 | def wrong(): # [too-many-branches] - | ^^^^^ PLR0912 -7 | """ Has too many branches. """ -8 | if 1: - | - +too_many_branches.py:8:5: PLR0912 Too many branches (13 > 12) + | + 7 | # pylint: disable=using-constant-test + 8 | def wrong(): # [too-many-branches] + | ^^^^^ PLR0912 + 9 | """ Has too many branches. """ +10 | if 1: + | +too_many_branches.py:76:5: PLR0912 Too many branches (13 > 12) + | +74 | pass +75 | +76 | def with_statement_wrong(): + | ^^^^^^^^^^^^^^^^^^^^ PLR0912 +77 | """statements inside the with statement should get counted""" +78 | with suppress(Exception): + |