Implement pylint's too-many-return-statements rule (PLR0911) (#2564)

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 | 🛠 |
+| PLR0911 | too-many-return-statements | Too many return statements ({returns}/{max_returns}) |  |
 | 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}) |  |
@@ -3775,6 +3776,23 @@ max-branches = 12
 
 ---
 
+#### [`max-returns`](#max-returns)
+
+Maximum number of return statements allowed for a function or method body (see `PLR0911`)
+
+**Default value**: `6`
+
+**Type**: `int`
+
+**Example usage**:
+
+```toml
+[tool.ruff.pylint]
+max-returns = 6
+```
+
+---
+
 #### [`max-statements`](#max-statements)
 
 Maximum number of statements allowed for a function or method body (see: `PLR0915`).

resources/test/fixtures/pylint/too_many_return_statements.py

@@ -0,0 +1,41 @@
+"""
+https://github.com/PyCQA/pylint/blob/69eca9b3f9856c3033957b769358803ee48e8e47/tests/functional/t/too/too_many_return_statements.py
+"""
+def stupid_function(arg): # [too-many-return-statements]
+    if arg == 1:
+        return 1
+    if arg == 2:
+        return 2
+    if arg == 3:
+        return 3
+    if arg == 4:
+        return 4
+    if arg == 5:
+        return 5
+    if arg == 6:
+        return 6
+    if arg == 7:
+        return 7
+    if arg == 8:
+        return 8
+    if arg == 9:
+        return 9
+    if arg == 10:
+        return 10
+    return None
+
+
+def many_yield(text):
+    """Not a problem"""
+    if text:
+        yield f"    line 1: {text}\n"
+        yield "    line 2\n"
+        yield "    line 3\n"
+        yield "    line 4\n"
+        yield "    line 5\n"
+    else:
+        yield "    line 6\n"
+        yield "    line 7\n"
+        yield "    line 8\n"
+        yield "    line 9\n"
+        yield "    line 10\n"

resources/test/fixtures/pylint/too_many_return_statements_params.py

@@ -0,0 +1,5 @@
+def f(x):  # Too many return statements (2/1)
+    if x == 1:
+        return
+    if x == 2:
+        return

ruff.schema.json

@@ -1199,6 +1199,15 @@
           "format": "uint",
           "minimum": 0.0
         },
+        "max-returns": {
+          "description": "Maximum number of return statements allowed for a function or method body (see `PLR0911`)",
+          "type": [
+            "integer",
+            "null"
+          ],
+          "format": "uint",
+          "minimum": 0.0
+        },
         "max-statements": {
           "description": "Maximum number of statements allowed for a function or method body (see: `PLR0915`).",
           "type": [
@@ -1678,6 +1687,7 @@
         "PLR0402",
         "PLR09",
         "PLR091",
+        "PLR0911",
         "PLR0912",
         "PLR0913",
         "PLR0915",

src/ast/helpers.rs

@@ -13,6 +13,8 @@ use rustpython_parser::token::StringKind;
 use smallvec::smallvec;
 
 use crate::ast::types::{Binding, BindingKind, CallPath, Range};
+use crate::ast::visitor;
+use crate::ast::visitor::Visitor;
 use crate::checkers::ast::Checker;
 use crate::source_code::{Generator, Indexer, Locator, Stylist};
 
@@ -664,6 +666,27 @@ pub fn to_call_path(target: &str) -> CallPath {
     }
 }
 
+/// A [`Visitor`] that collects all return statements in a function or method.
+#[derive(Default)]
+pub struct ReturnStatementVisitor<'a> {
+    pub returns: Vec<Option<&'a Expr>>,
+}
+
+impl<'a, 'b> Visitor<'b> for ReturnStatementVisitor<'a>
+where
+    'b: 'a,
+{
+    fn visit_stmt(&mut self, stmt: &'b Stmt) {
+        match &stmt.node {
+            StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
+                // Don't recurse.
+            }
+            StmtKind::Return { value } => self.returns.push(value.as_ref().map(|expr| &**expr)),
+            _ => visitor::walk_stmt(self, stmt),
+        }
+    }
+}
+
 /// Convert a location within a file (relative to `base`) to an absolute
 /// position.
 pub fn to_absolute(relative: Location, base: Location) -> Location {

src/checkers/ast.rs

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

src/registry.rs

@@ -92,6 +92,7 @@ ruff_macros::define_rule_mapping!(
     PLR2004 => rules::pylint::rules::MagicValueComparison,
     PLW0120 => rules::pylint::rules::UselessElseOnLoop,
     PLW0602 => rules::pylint::rules::GlobalVariableNotAssigned,
+    PLR0911 => rules::pylint::rules::TooManyReturnStatements,
     PLR0913 => rules::pylint::rules::TooManyArguments,
     PLR0912 => rules::pylint::rules::TooManyBranches,
     PLR0915 => rules::pylint::rules::TooManyStatements,

src/rules/flake8_annotations/rules.rs

@@ -1,11 +1,12 @@
 use log::error;
-use rustpython_ast::{Constant, Expr, ExprKind, Stmt, StmtKind};
+use rustpython_ast::{Constant, Expr, ExprKind, Stmt};
 
-use super::fixes;
-use super::helpers::match_function_def;
+use ruff_macros::derive_message_formats;
+
+use crate::ast::helpers::ReturnStatementVisitor;
 use crate::ast::types::Range;
 use crate::ast::visitor::Visitor;
-use crate::ast::{cast, helpers, visitor};
+use crate::ast::{cast, helpers};
 use crate::checkers::ast::Checker;
 use crate::define_violation;
 use crate::docstrings::definition::{Definition, DefinitionKind};
@@ -13,7 +14,9 @@ use crate::registry::{Diagnostic, Rule};
 use crate::violation::{AlwaysAutofixableViolation, Violation};
 use crate::visibility;
 use crate::visibility::Visibility;
-use ruff_macros::derive_message_formats;
+
+use super::fixes;
+use super::helpers::match_function_def;
 
 define_violation!(
     pub struct MissingTypeFunctionArgument {
@@ -162,26 +165,6 @@ impl Violation for DynamicallyTypedExpression {
     }
 }
 
-#[derive(Default)]
-struct ReturnStatementVisitor<'a> {
-    returns: Vec<Option<&'a Expr>>,
-}
-
-impl<'a, 'b> Visitor<'b> for ReturnStatementVisitor<'a>
-where
-    'b: 'a,
-{
-    fn visit_stmt(&mut self, stmt: &'b Stmt) {
-        match &stmt.node {
-            StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
-                // Don't recurse.
-            }
-            StmtKind::Return { value } => self.returns.push(value.as_ref().map(|expr| &**expr)),
-            _ => visitor::walk_stmt(self, stmt),
-        }
-    }
-}
-
 fn is_none_returning(body: &[Stmt]) -> bool {
     let mut visitor = ReturnStatementVisitor::default();
     for stmt in body {

src/rules/pylint/mod.rs

@@ -36,6 +36,7 @@ mod tests {
     #[test_case(Rule::GlobalVariableNotAssigned, Path::new("global_variable_not_assigned.py"); "PLW0602")]
     #[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::TooManyReturnStatements, Path::new("too_many_return_statements.py"); "PLR0911")]
     #[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")]
@@ -125,4 +126,20 @@ mod tests {
         assert_yaml_snapshot!(diagnostics);
         Ok(())
     }
+
+    #[test]
+    fn max_return_statements() -> Result<()> {
+        let diagnostics = test_path(
+            Path::new("pylint/too_many_return_statements_params.py"),
+            &Settings {
+                pylint: pylint::settings::Settings {
+                    max_returns: 1,
+                    ..pylint::settings::Settings::default()
+                },
+                ..Settings::for_rules(vec![Rule::TooManyReturnStatements])
+            },
+        )?;
+        assert_yaml_snapshot!(diagnostics);
+        Ok(())
+    }
 }

src/rules/pylint/rules/mod.rs

@@ -10,6 +10,7 @@ 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_return_statements::{too_many_return_statements, TooManyReturnStatements};
 pub use too_many_statements::{too_many_statements, TooManyStatements};
 pub use unnecessary_direct_lambda_call::{
     unnecessary_direct_lambda_call, UnnecessaryDirectLambdaCall,
@@ -33,6 +34,7 @@ mod nonlocal_without_binding;
 mod property_with_parameters;
 mod too_many_arguments;
 mod too_many_branches;
+mod too_many_return_statements;
 mod too_many_statements;
 mod unnecessary_direct_lambda_call;
 mod use_from_import;

src/rules/pylint/rules/too_many_return_statements.rs

@@ -0,0 +1,216 @@
+use crate::ast::helpers::{identifier_range, ReturnStatementVisitor};
+use crate::ast::visitor::Visitor;
+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::Stmt;
+
+define_violation!(
+    pub struct TooManyReturnStatements {
+        pub returns: usize,
+        pub max_returns: usize,
+    }
+);
+impl Violation for TooManyReturnStatements {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let TooManyReturnStatements {
+            returns,
+            max_returns,
+        } = self;
+        format!("Too many return statements ({returns}/{max_returns})")
+    }
+}
+
+/// Count the number of return statements in a function or method body.
+fn num_returns(body: &[Stmt]) -> usize {
+    let mut visitor = ReturnStatementVisitor::default();
+    visitor.visit_body(body);
+    visitor.returns.len()
+}
+
+/// PLR0911
+pub fn too_many_return_statements(
+    stmt: &Stmt,
+    body: &[Stmt],
+    max_returns: usize,
+    locator: &Locator,
+) -> Option<Diagnostic> {
+    let returns = num_returns(body);
+    if returns > max_returns {
+        Some(Diagnostic::new(
+            TooManyReturnStatements {
+                returns,
+                max_returns,
+            },
+            identifier_range(stmt, locator),
+        ))
+    } else {
+        None
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use rustpython_parser::parser;
+
+    use super::num_returns;
+
+    fn test_helper(source: &str, expected: usize) -> Result<()> {
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(num_returns(&stmts), expected);
+        Ok(())
+    }
+
+    #[test]
+    fn if_() -> Result<()> {
+        let source = r#"
+x = 1
+if x == 1:  # 9
+    return
+if x == 2:
+    return
+if x == 3:
+    return
+if x == 4:
+    return
+if x == 5:
+    return
+if x == 6:
+    return
+if x == 7:
+    return
+if x == 8:
+    return
+if x == 9:
+    return
+"#;
+
+        test_helper(source, 9)?;
+        Ok(())
+    }
+
+    #[test]
+    fn for_else() -> Result<()> {
+        let source = r#"
+for _i in range(10):
+    return
+else:
+    return
+"#;
+
+        test_helper(source, 2)?;
+        Ok(())
+    }
+
+    #[test]
+    fn async_for_else() -> Result<()> {
+        let source = r#"
+async for _i in range(10):
+    return
+else:
+    return
+"#;
+
+        test_helper(source, 2)?;
+        Ok(())
+    }
+
+    #[test]
+    fn nested_def_ignored() -> Result<()> {
+        let source = r#"
+def f():
+    return
+
+x = 1
+if x == 1:
+    print()
+else:
+    print()
+"#;
+
+        test_helper(source, 0)?;
+        Ok(())
+    }
+
+    #[test]
+    fn while_nested_if() -> Result<()> {
+        let source = r#"
+x = 1
+while x < 10:
+    print()
+    if x == 3:
+        return
+    x += 1
+return
+"#;
+
+        test_helper(source, 2)?;
+        Ok(())
+    }
+
+    #[test]
+    fn with_if() -> Result<()> {
+        let source = r#"
+with a as f:
+    return
+    if f == 1:
+        return
+"#;
+
+        test_helper(source, 2)?;
+        Ok(())
+    }
+
+    #[test]
+    fn async_with_if() -> Result<()> {
+        let source = r#"
+async with a as f:
+    return
+    if f == 1:
+        return
+"#;
+
+        test_helper(source, 2)?;
+        Ok(())
+    }
+
+    #[test]
+    fn try_except_except_else_finally() -> Result<()> {
+        let source = r#"
+try:
+    print()
+    return
+except ValueError:
+    return
+except Exception:
+    return
+else:
+    return
+finally:
+    return
+"#;
+
+        test_helper(source, 5)?;
+        Ok(())
+    }
+
+    #[test]
+    fn class_def_ignored() -> Result<()> {
+        let source = r#"
+class A:
+    def f(self):
+        return
+
+    def g(self):
+        return
+"#;
+
+        test_helper(source, 0)?;
+        Ok(())
+    }
+}

src/rules/pylint/settings.rs

@@ -56,6 +56,9 @@ pub struct Options {
     #[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"6", value_type = "int", example = r"max-returns = 6")]
+    /// Maximum number of return statements allowed for a function or method body (see `PLR0911`)
+    pub max_returns: Option<usize>,
     #[option(default = r"5", value_type = "int", example = r"max-args = 5")]
     /// Maximum number of arguments allowed for a function or method definition (see: `PLR0913`).
     pub max_args: Option<usize>,
@@ -68,6 +71,7 @@ pub struct Options {
 pub struct Settings {
     pub allow_magic_value_types: Vec<ConstantType>,
     pub max_args: usize,
+    pub max_returns: usize,
     pub max_branches: usize,
     pub max_statements: usize,
 }
@@ -77,6 +81,7 @@ impl Default for Settings {
         Self {
             allow_magic_value_types: vec![ConstantType::Str, ConstantType::Bytes],
             max_args: 5,
+            max_returns: 6,
             max_branches: 12,
             max_statements: 50,
         }
@@ -91,6 +96,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_returns: options.max_returns.unwrap_or(defaults.max_returns),
             max_branches: options.max_branches.unwrap_or(defaults.max_branches),
             max_statements: options.max_statements.unwrap_or(defaults.max_statements),
         }
@@ -102,6 +108,7 @@ impl From<Settings> for Options {
         Self {
             allow_magic_value_types: Some(settings.allow_magic_value_types),
             max_args: Some(settings.max_args),
+            max_returns: Some(settings.max_returns),
             max_branches: Some(settings.max_branches),
             max_statements: Some(settings.max_statements),
         }

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

@@ -0,0 +1,17 @@
+---
+source: src/rules/pylint/mod.rs
+expression: diagnostics
+---
+- kind:
+    TooManyReturnStatements:
+      returns: 11
+      max_returns: 6
+  location:
+    row: 4
+    column: 4
+  end_location:
+    row: 4
+    column: 19
+  fix: ~
+  parent: ~
+

src/rules/pylint/snapshots/ruff__rules__pylint__tests__max_return_statements.snap

@@ -0,0 +1,17 @@
+---
+source: src/rules/pylint/mod.rs
+expression: diagnostics
+---
+- kind:
+    TooManyReturnStatements:
+      returns: 2
+      max_returns: 1
+  location:
+    row: 1
+    column: 4
+  end_location:
+    row: 1
+    column: 5
+  fix: ~
+  parent: ~
+