diff --git a/resources/test/fixtures/flake8_return/RET503.py b/resources/test/fixtures/flake8_return/RET503.py index 13886ce55a..0cf189612f 100644 --- a/resources/test/fixtures/flake8_return/RET503.py +++ b/resources/test/fixtures/flake8_return/RET503.py @@ -114,3 +114,12 @@ def bar3(x, y, z): else: return z return None + + +def prompts(self, foo): + if not foo: + return [] + + for x in foo: + yield x + yield x + 1 diff --git a/src/rules/flake8_return/rules.rs b/src/rules/flake8_return/rules.rs index cba99e437d..f9e289a06e 100644 --- a/src/rules/flake8_return/rules.rs +++ b/src/rules/flake8_return/rules.rs @@ -325,6 +325,11 @@ pub fn function(checker: &mut Checker, body: &[Stmt]) { visitor.stack }; + // Avoid false positives for generators. + if !stack.yields.is_empty() { + return; + } + if checker.settings.rules.enabled(&Rule::SuperfluousElseReturn) || checker.settings.rules.enabled(&Rule::SuperfluousElseRaise) || checker diff --git a/src/rules/flake8_return/visitor.rs b/src/rules/flake8_return/visitor.rs index 10a05f59ff..4e58d06547 100644 --- a/src/rules/flake8_return/visitor.rs +++ b/src/rules/flake8_return/visitor.rs @@ -7,6 +7,7 @@ use crate::ast::visitor::Visitor; #[derive(Default)] pub struct Stack<'a> { pub returns: Vec<(&'a Stmt, Option<&'a Expr>)>, + pub yields: Vec<&'a Expr>, pub ifs: Vec<&'a Stmt>, pub elifs: Vec<&'a Stmt>, pub refs: FxHashMap<&'a str, Vec>, @@ -116,7 +117,6 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> { .push((stmt.location, stmt.end_location.unwrap())); visitor::walk_stmt(self, stmt); } - _ => { visitor::walk_stmt(self, stmt); } @@ -143,8 +143,8 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> { .or_insert_with(Vec::new) .push(expr.location); } - ExprKind::JoinedStr { .. } => { - visitor::walk_expr(self, expr); + ExprKind::YieldFrom { .. } | ExprKind::Yield { .. } => { + self.stack.yields.push(expr); } _ => visitor::walk_expr(self, expr), }