Fix false-positive in RET504 when referencing globals

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

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),

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<Location>>,
+    pub non_locals: FxHashSet<&'a str>,
     pub assigns: FxHashMap<&'a str, Vec<Location>>,
     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.
             }