From ced55084db537a8591c1dba9d625ed1d33acc235 Mon Sep 17 00:00:00 2001 From: Chris Chan Date: Sat, 4 Feb 2023 21:56:36 +0000 Subject: [PATCH] Implement pylint's `too-many-return-statements` rule (`PLR0911`) (#2564) --- README.md | 18 ++ .../pylint/too_many_return_statements.py | 41 ++++ .../too_many_return_statements_params.py | 5 + ruff.schema.json | 10 + src/ast/helpers.rs | 23 ++ src/checkers/ast.rs | 11 + src/registry.rs | 1 + src/rules/flake8_annotations/rules.rs | 33 +-- src/rules/pylint/mod.rs | 17 ++ src/rules/pylint/rules/mod.rs | 2 + .../rules/too_many_return_statements.rs | 216 ++++++++++++++++++ src/rules/pylint/settings.rs | 7 + ...PLR0911_too_many_return_statements.py.snap | 17 ++ ..._pylint__tests__max_return_statements.snap | 17 ++ 14 files changed, 393 insertions(+), 25 deletions(-) create mode 100644 resources/test/fixtures/pylint/too_many_return_statements.py create mode 100644 resources/test/fixtures/pylint/too_many_return_statements_params.py create mode 100644 src/rules/pylint/rules/too_many_return_statements.rs create mode 100644 src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR0911_too_many_return_statements.py.snap create mode 100644 src/rules/pylint/snapshots/ruff__rules__pylint__tests__max_return_statements.snap diff --git a/README.md b/README.md index 233ae3b6f0..485035dd61 100644 --- a/README.md +++ b/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`). diff --git a/resources/test/fixtures/pylint/too_many_return_statements.py b/resources/test/fixtures/pylint/too_many_return_statements.py new file mode 100644 index 0000000000..a2241be7e1 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/resources/test/fixtures/pylint/too_many_return_statements_params.py b/resources/test/fixtures/pylint/too_many_return_statements_params.py new file mode 100644 index 0000000000..abbd52101a --- /dev/null +++ b/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 diff --git a/ruff.schema.json b/ruff.schema.json index d3b222eb00..b667e0594f 100644 --- a/ruff.schema.json +++ b/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", diff --git a/src/ast/helpers.rs b/src/ast/helpers.rs index 24c59bab25..a8c4d88596 100644 --- a/src/ast/helpers.rs +++ b/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>, +} + +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 { diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs index 8597ef9ef8..933e133f43 100644 --- a/src/checkers/ast.rs +++ b/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, diff --git a/src/registry.rs b/src/registry.rs index 9dfc0a195d..5c19d40b63 100644 --- a/src/registry.rs +++ b/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, diff --git a/src/rules/flake8_annotations/rules.rs b/src/rules/flake8_annotations/rules.rs index 764d6c22dc..1cbe6c6ed4 100644 --- a/src/rules/flake8_annotations/rules.rs +++ b/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>, -} - -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 { diff --git a/src/rules/pylint/mod.rs b/src/rules/pylint/mod.rs index 73cc95b38c..6b421cb84b 100644 --- a/src/rules/pylint/mod.rs +++ b/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(()) + } } diff --git a/src/rules/pylint/rules/mod.rs b/src/rules/pylint/rules/mod.rs index afa8bde1be..07dfb9cefd 100644 --- a/src/rules/pylint/rules/mod.rs +++ b/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; diff --git a/src/rules/pylint/rules/too_many_return_statements.rs b/src/rules/pylint/rules/too_many_return_statements.rs new file mode 100644 index 0000000000..77def4e004 --- /dev/null +++ b/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 { + 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, "")?; + 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(()) + } +} diff --git a/src/rules/pylint/settings.rs b/src/rules/pylint/settings.rs index c2b775378a..6497e93437 100644 --- a/src/rules/pylint/settings.rs +++ b/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, + #[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, #[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, @@ -68,6 +71,7 @@ pub struct Options { pub struct Settings { pub allow_magic_value_types: Vec, 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 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 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), } diff --git a/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR0911_too_many_return_statements.py.snap b/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR0911_too_many_return_statements.py.snap new file mode 100644 index 0000000000..2628e5bdcc --- /dev/null +++ b/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: ~ + diff --git a/src/rules/pylint/snapshots/ruff__rules__pylint__tests__max_return_statements.snap b/src/rules/pylint/snapshots/ruff__rules__pylint__tests__max_return_statements.snap new file mode 100644 index 0000000000..0619b609b0 --- /dev/null +++ b/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: ~ +