Implement pylint's too-many-branches rule (PLR0912) (#2550)

README.md

@@ -1343,6 +1343,7 @@ For more, see [Pylint](https://pypi.org/project/pylint/) on PyPI.
 | PLR0133 | comparison-of-constant | Two constants compared in a comparison, consider replacing `{left_constant} {op} {right_constant}` |  |
 | PLR0206 | property-with-parameters | Cannot have defined parameters for properties |  |
 | PLR0402 | consider-using-from-import | Use `from {module} import {name}` in lieu of alias | 🛠 |
+| PLR0912 | too-many-branches | Too many branches ({branches}/{max_branches}) |  |
 | PLR0913 | too-many-arguments | Too many arguments to function call ({c_args}/{max_args}) |  |
 | PLR0915 | too-many-statements | Too many statements ({statements}/{max_statements}) |  |
 | PLR1701 | consider-merging-isinstance | Merge these isinstance calls: `isinstance({obj}, ({types}))` |  |
@@ -3742,7 +3743,7 @@ allow-magic-value-types = ["int"]
 
 #### [`max-args`](#max-args)
 
-Maximum number of arguments allowed for a function definition (see: `PLR0913`).
+Maximum number of arguments allowed for a function or method definition (see: `PLR0913`).
 
 **Default value**: `5`
 
@@ -3757,9 +3758,26 @@ max-args = 5
 
 ---
 
+#### [`max-branches`](#max-branches)
+
+Maximum number of branches allowed for a function or method body (see: `PLR0912`).
+
+**Default value**: `12`
+
+**Type**: `int`
+
+**Example usage**:
+
+```toml
+[tool.ruff.pylint]
+max-branches = 12
+```
+
+---
+
 #### [`max-statements`](#max-statements)
 
-Maximum number of statements allowed for a method or a statement (see: `PLR0915`).
+Maximum number of statements allowed for a function or method body (see: `PLR0915`).
 
 **Default value**: `50`
 

resources/test/fixtures/pylint/too_many_branches.py

@@ -0,0 +1,72 @@
+"""
+Test for too many branches.
+Taken from the pylint source 2023-02-03
+"""
+# pylint: disable=using-constant-test
+def wrong(): # [too-many-branches]
+    """ Has too many branches. """
+    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
+
+def good():
+    """ Too many branches only if we take
+    into consideration the nested functions.
+    """
+    def nested_1():
+        """ empty """
+        if 1:
+            pass
+        elif 2:
+            pass
+        elif 3:
+            pass
+        elif 4:
+            pass
+
+    nested_1()
+    try:
+        pass
+    finally:
+        pass
+    try:
+        pass
+    finally:
+        pass
+    if 1:
+        pass
+    elif 2:
+        pass
+    elif 3:
+        pass
+    elif 4:
+        pass
+    elif 5:
+        pass
+    elif 6:
+        pass
+    elif 7:
+        pass

resources/test/fixtures/pylint/too_many_branches_params.py

@@ -0,0 +1,19 @@
+def f(x):
+    if x:
+        pass
+
+
+def g(x):
+    if x:
+        pass
+    else:
+        pass
+
+    def h(x):
+        return
+
+    def i(x):
+        if x:
+            pass
+        else:
+            pass

ruff.schema.json

@@ -1182,7 +1182,16 @@
           }
         },
         "max-args": {
-          "description": "Maximum number of arguments allowed for a function definition (see: `PLR0913`).",
+          "description": "Maximum number of arguments allowed for a function or method definition (see: `PLR0913`).",
+          "type": [
+            "integer",
+            "null"
+          ],
+          "format": "uint",
+          "minimum": 0.0
+        },
+        "max-branches": {
+          "description": "Maximum number of branches allowed for a function or method body (see: `PLR0912`).",
           "type": [
             "integer",
             "null"
@@ -1191,7 +1200,7 @@
           "minimum": 0.0
         },
         "max-statements": {
-          "description": "Maximum number of statements allowed for a method or a statement (see: `PLR0915`).",
+          "description": "Maximum number of statements allowed for a function or method body (see: `PLR0915`).",
           "type": [
             "integer",
             "null"
@@ -1669,6 +1678,7 @@
         "PLR0402",
         "PLR09",
         "PLR091",
+        "PLR0912",
         "PLR0913",
         "PLR0915",
         "PLR1",

src/checkers/ast.rs

@@ -592,6 +592,17 @@ where
                     pylint::rules::too_many_arguments(self, args, stmt);
                 }
 
+                if self.settings.rules.enabled(&Rule::TooManyBranches) {
+                    if let Some(diagnostic) = pylint::rules::too_many_branches(
+                        stmt,
+                        body,
+                        self.settings.pylint.max_branches,
+                        self.locator,
+                    ) {
+                        self.diagnostics.push(diagnostic);
+                    }
+                }
+
                 if self.settings.rules.enabled(&Rule::TooManyStatements) {
                     if let Some(diagnostic) = pylint::rules::too_many_statements(
                         stmt,

src/registry.rs

@@ -93,6 +93,7 @@ ruff_macros::define_rule_mapping!(
     PLW0120 => rules::pylint::rules::UselessElseOnLoop,
     PLW0602 => rules::pylint::rules::GlobalVariableNotAssigned,
     PLR0913 => rules::pylint::rules::TooManyArguments,
+    PLR0912 => rules::pylint::rules::TooManyBranches,
     PLR0915 => rules::pylint::rules::TooManyStatements,
     // flake8-builtins
     A001 => rules::flake8_builtins::rules::BuiltinVariableShadowing,

src/rules/pylint/mod.rs

@@ -37,6 +37,7 @@ mod tests {
     #[test_case(Rule::InvalidAllFormat, Path::new("invalid_all_format.py"); "PLE0605")]
     #[test_case(Rule::InvalidAllObject, Path::new("invalid_all_object.py"); "PLE0604")]
     #[test_case(Rule::TooManyArguments, Path::new("too_many_arguments.py"); "PLR0913")]
+    #[test_case(Rule::TooManyBranches, Path::new("too_many_branches.py"); "PLR0912")]
     #[test_case(Rule::TooManyStatements, Path::new("too_many_statements.py"); "PLR0915")]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.code(), path.to_string_lossy());
@@ -93,6 +94,22 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn max_branches() -> Result<()> {
+        let diagnostics = test_path(
+            Path::new("pylint/too_many_branches_params.py"),
+            &Settings {
+                pylint: pylint::settings::Settings {
+                    max_branches: 1,
+                    ..pylint::settings::Settings::default()
+                },
+                ..Settings::for_rules(vec![Rule::TooManyBranches])
+            },
+        )?;
+        assert_yaml_snapshot!(diagnostics);
+        Ok(())
+    }
+
     #[test]
     fn max_statements() -> Result<()> {
         let diagnostics = test_path(

src/rules/pylint/rules/mod.rs

@@ -9,6 +9,7 @@ pub use merge_isinstance::{merge_isinstance, ConsiderMergingIsinstance};
 pub use nonlocal_without_binding::NonlocalWithoutBinding;
 pub use property_with_parameters::{property_with_parameters, PropertyWithParameters};
 pub use too_many_arguments::{too_many_arguments, TooManyArguments};
+pub use too_many_branches::{too_many_branches, TooManyBranches};
 pub use too_many_statements::{too_many_statements, TooManyStatements};
 pub use unnecessary_direct_lambda_call::{
     unnecessary_direct_lambda_call, UnnecessaryDirectLambdaCall,
@@ -31,6 +32,7 @@ mod merge_isinstance;
 mod nonlocal_without_binding;
 mod property_with_parameters;
 mod too_many_arguments;
+mod too_many_branches;
 mod too_many_statements;
 mod unnecessary_direct_lambda_call;
 mod use_from_import;

src/rules/pylint/rules/too_many_branches.rs

@@ -0,0 +1,204 @@
+use crate::ast::helpers::identifier_range;
+use crate::define_violation;
+use crate::registry::Diagnostic;
+use crate::source_code::Locator;
+use crate::violation::Violation;
+
+use ruff_macros::derive_message_formats;
+use rustpython_ast::{ExcepthandlerKind, Stmt, StmtKind};
+
+define_violation!(
+    pub struct TooManyBranches {
+        pub branches: usize,
+        pub max_branches: usize,
+    }
+);
+
+impl Violation for TooManyBranches {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let TooManyBranches {
+            branches,
+            max_branches,
+        } = self;
+        format!("Too many branches ({branches}/{max_branches})")
+    }
+}
+
+fn num_branches(stmts: &[Stmt]) -> usize {
+    stmts
+        .iter()
+        .map(|stmt| {
+            // TODO(charlie): Account for pattern match statement.
+            match &stmt.node {
+                StmtKind::If { body, orelse, .. } => {
+                    1 + num_branches(body)
+                        + (if let Some(stmt) = orelse.first() {
+                            // `elif:` and `else: if:` have the same AST representation.
+                            // Avoid treating `elif:` as two statements.
+                            usize::from(!matches!(stmt.node, StmtKind::If { .. }))
+                        } else {
+                            0
+                        })
+                        + num_branches(orelse)
+                }
+                StmtKind::For { body, orelse, .. }
+                | StmtKind::AsyncFor { body, orelse, .. }
+                | StmtKind::While { body, orelse, .. } => {
+                    1 + num_branches(body)
+                        + (if orelse.is_empty() {
+                            0
+                        } else {
+                            1 + num_branches(orelse)
+                        })
+                }
+                StmtKind::Try {
+                    body,
+                    handlers,
+                    orelse,
+                    finalbody,
+                } => {
+                    1 + num_branches(body)
+                        + (if orelse.is_empty() {
+                            0
+                        } else {
+                            1 + num_branches(orelse)
+                        })
+                        + (if finalbody.is_empty() {
+                            0
+                        } else {
+                            1 + num_branches(finalbody)
+                        })
+                        + handlers
+                            .iter()
+                            .map(|handler| {
+                                1 + {
+                                    let ExcepthandlerKind::ExceptHandler { body, .. } =
+                                        &handler.node;
+                                    num_branches(body)
+                                }
+                            })
+                            .sum::<usize>()
+                }
+                _ => 0,
+            }
+        })
+        .sum()
+}
+
+/// PLR0912
+pub fn too_many_branches(
+    stmt: &Stmt,
+    body: &[Stmt],
+    max_branches: usize,
+    locator: &Locator,
+) -> Option<Diagnostic> {
+    let branches = num_branches(body);
+    if branches > max_branches {
+        Some(Diagnostic::new(
+            TooManyBranches {
+                branches,
+                max_branches,
+            },
+            identifier_range(stmt, locator),
+        ))
+    } else {
+        None
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use rustpython_parser::parser;
+
+    use super::num_branches;
+
+    fn test_helper(source: &str, expected_num_branches: usize) -> Result<()> {
+        let branches = parser::parse_program(source, "<filename>")?;
+        assert_eq!(num_branches(&branches), expected_num_branches);
+        Ok(())
+    }
+
+    #[test]
+    fn if_else_nested_if_else() -> Result<()> {
+        let source: &str = r#"
+if x == 0:  # 3
+    return
+else:
+    if x == 1:
+        pass
+    else:
+        pass
+"#;
+
+        test_helper(source, 3)?;
+        Ok(())
+    }
+
+    #[test]
+    fn for_else() -> Result<()> {
+        let source: &str = r#"
+for _ in range(x):  # 2
+    pass
+else:
+    pass
+"#;
+
+        test_helper(source, 2)?;
+        Ok(())
+    }
+
+    #[test]
+    fn while_if_else_if() -> Result<()> {
+        let source: &str = r#"
+while x < 1:  # 4
+    if x:
+        pass
+else:
+    if x:
+        pass
+"#;
+
+        test_helper(source, 4)?;
+        Ok(())
+    }
+
+    #[test]
+    fn nested_def() -> Result<()> {
+        let source: &str = r#"
+if x:  # 2
+    pass
+else:
+    pass
+
+def g(x):
+    if x:
+        pass
+
+return 1
+"#;
+
+        test_helper(source, 2)?;
+        Ok(())
+    }
+
+    #[test]
+    fn try_except_except_else_finally() -> Result<()> {
+        let source: &str = r#"
+try:
+    pass
+except:
+    pass
+except:
+    pass
+else:
+    pass
+finally:
+    pass
+"#;
+
+        test_helper(source, 5)?;
+        Ok(())
+    }
+}

src/rules/pylint/settings.rs

@@ -53,11 +53,14 @@ pub struct Options {
     )]
     /// Constant types to ignore when used as "magic values" (see: `PLR2004`).
     pub allow_magic_value_types: Option<Vec<ConstantType>>,
+    #[option(default = r"12", value_type = "int", example = r"max-branches = 12")]
+    /// Maximum number of branches allowed for a function or method body (see: `PLR0912`).
+    pub max_branches: Option<usize>,
     #[option(default = r"5", value_type = "int", example = r"max-args = 5")]
-    /// Maximum number of arguments allowed for a function definition (see: `PLR0913`).
+    /// Maximum number of arguments allowed for a function or method definition (see: `PLR0913`).
     pub max_args: Option<usize>,
     #[option(default = r"50", value_type = "int", example = r"max-statements = 50")]
-    /// Maximum number of statements allowed for a method or a statement (see: `PLR0915`).
+    /// Maximum number of statements allowed for a function or method body (see: `PLR0915`).
     pub max_statements: Option<usize>,
 }
 
@@ -65,6 +68,7 @@ pub struct Options {
 pub struct Settings {
     pub allow_magic_value_types: Vec<ConstantType>,
     pub max_args: usize,
+    pub max_branches: usize,
     pub max_statements: usize,
 }
 
@@ -73,6 +77,7 @@ impl Default for Settings {
         Self {
             allow_magic_value_types: vec![ConstantType::Str, ConstantType::Bytes],
             max_args: 5,
+            max_branches: 12,
             max_statements: 50,
         }
     }
@@ -86,6 +91,7 @@ impl From<Options> for Settings {
                 .allow_magic_value_types
                 .unwrap_or(defaults.allow_magic_value_types),
             max_args: options.max_args.unwrap_or(defaults.max_args),
+            max_branches: options.max_branches.unwrap_or(defaults.max_branches),
             max_statements: options.max_statements.unwrap_or(defaults.max_statements),
         }
     }
@@ -96,6 +102,7 @@ impl From<Settings> for Options {
         Self {
             allow_magic_value_types: Some(settings.allow_magic_value_types),
             max_args: Some(settings.max_args),
+            max_branches: Some(settings.max_branches),
             max_statements: Some(settings.max_statements),
         }
     }

src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR0912_too_many_branches.py.snap

@@ -0,0 +1,17 @@
+---
+source: src/rules/pylint/mod.rs
+expression: diagnostics
+---
+- kind:
+    TooManyBranches:
+      branches: 13
+      max_branches: 12
+  location:
+    row: 6
+    column: 4
+  end_location:
+    row: 6
+    column: 9
+  fix: ~
+  parent: ~
+

src/rules/pylint/snapshots/ruff__rules__pylint__tests__max_branches.snap

@@ -0,0 +1,29 @@
+---
+source: src/rules/pylint/mod.rs
+expression: diagnostics
+---
+- kind:
+    TooManyBranches:
+      branches: 2
+      max_branches: 1
+  location:
+    row: 6
+    column: 4
+  end_location:
+    row: 6
+    column: 5
+  fix: ~
+  parent: ~
+- kind:
+    TooManyBranches:
+      branches: 2
+      max_branches: 1
+  location:
+    row: 15
+    column: 8
+  end_location:
+    row: 15
+    column: 9
+  fix: ~
+  parent: ~
+