diff --git a/resources/test/fixtures/flake8_return/RET504.py b/resources/test/fixtures/flake8_return/RET504.py index 76ff4d304b1ec..6f0e6f9b9c0fe 100644 --- a/resources/test/fixtures/flake8_return/RET504.py +++ b/resources/test/fixtures/flake8_return/RET504.py @@ -237,3 +237,15 @@ def close(self): any_failed = True report(traceback.format_exc()) return any_failed + +def global_assignment(): + global X + X = 1 + return X + +def nonlocal_assignment(): + X = 1 + def inner(): + nonlocal X + X = 1 + return X diff --git a/src/flake8_return/plugins.rs b/src/flake8_return/plugins.rs index 03b6ceb54a46f..c2f37edc19610 100644 --- a/src/flake8_return/plugins.rs +++ b/src/flake8_return/plugins.rs @@ -203,6 +203,10 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack, expr: &Expr) { return; } + if stack.non_locals.contains(id.as_str()) { + return; + } + checker.add_check(Check::new( CheckKind::UnnecessaryAssign, Range::from_located(expr), diff --git a/src/flake8_return/visitor.rs b/src/flake8_return/visitor.rs index b999a9b900a3f..e76cbe4605cf1 100644 --- a/src/flake8_return/visitor.rs +++ b/src/flake8_return/visitor.rs @@ -1,4 +1,4 @@ -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; use rustpython_ast::{Expr, ExprKind, Location, Stmt, StmtKind}; use crate::ast::visitor; @@ -10,6 +10,7 @@ pub struct Stack<'a> { pub ifs: Vec<&'a Stmt>, pub elifs: Vec<&'a Stmt>, pub refs: FxHashMap<&'a str, Vec>, + pub non_locals: FxHashSet<&'a str>, pub assigns: FxHashMap<&'a str, Vec>, pub loops: Vec<(Location, Location)>, pub tries: Vec<(Location, Location)>, @@ -48,6 +49,11 @@ impl<'a> ReturnVisitor<'a> { impl<'a> Visitor<'a> for ReturnVisitor<'a> { fn visit_stmt(&mut self, stmt: &'a Stmt) { match &stmt.node { + StmtKind::Global { names } | StmtKind::Nonlocal { names } => { + self.stack + .non_locals + .extend(names.iter().map(std::string::String::as_str)); + } StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => { // Don't recurse. }