Fix false-positive in RET504 when referencing globals

Closes #1310
astral-sh · Dec 24, 2022 · 4007ddb · 4007ddb
1 parent 4ded155
commit 4007ddb
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 13 deletions.
diff --git a/resources/test/fixtures/flake8_return/RET504.py b/resources/test/fixtures/flake8_return/RET504.py
@@ -237,3 +237,8 @@ def close(self):
             any_failed = True
             report(traceback.format_exc())
     return any_failed
+
+def global_assignment():
+    global X
+    X = 1
+    return X
diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs
@@ -466,18 +466,6 @@ where
                     flake8_bugbear::plugins::cached_instance_method(self, decorator_list);
                 }
 
-                if self.settings.enabled.contains(&CheckCode::RET501)
-                    || self.settings.enabled.contains(&CheckCode::RET502)
-                    || self.settings.enabled.contains(&CheckCode::RET503)
-                    || self.settings.enabled.contains(&CheckCode::RET504)
-                    || self.settings.enabled.contains(&CheckCode::RET505)
-                    || self.settings.enabled.contains(&CheckCode::RET506)
-                    || self.settings.enabled.contains(&CheckCode::RET507)
-                    || self.settings.enabled.contains(&CheckCode::RET508)
-                {
-                    flake8_return::plugins::function(self, body);
-                }
-
                 if self.settings.enabled.contains(&CheckCode::C901) {
                     if let Some(check) = mccabe::checks::function_is_too_complex(
                         stmt,
@@ -1269,6 +1257,18 @@ where
                     globals,
                 })));
 
+                if self.settings.enabled.contains(&CheckCode::RET501)
+                    || self.settings.enabled.contains(&CheckCode::RET502)
+                    || self.settings.enabled.contains(&CheckCode::RET503)
+                    || self.settings.enabled.contains(&CheckCode::RET504)
+                    || self.settings.enabled.contains(&CheckCode::RET505)
+                    || self.settings.enabled.contains(&CheckCode::RET506)
+                    || self.settings.enabled.contains(&CheckCode::RET507)
+                    || self.settings.enabled.contains(&CheckCode::RET508)
+                {
+                    flake8_return::plugins::function(self, body);
+                }
+
                 self.deferred_functions.push((
                     stmt,
                     (self.scope_stack.clone(), self.parents.clone()),

diff --git a/src/flake8_return/plugins.rs b/src/flake8_return/plugins.rs
@@ -1,7 +1,8 @@
 use itertools::Itertools;
+use log::error;
 use rustpython_ast::{Constant, Expr, ExprKind, Location, Stmt, StmtKind};
 
-use crate::ast::types::Range;
+use crate::ast::types::{Range, ScopeKind};
 use crate::ast::visitor::Visitor;
 use crate::ast::whitespace::indentation;
 use crate::autofix::Fix;
@@ -203,6 +204,14 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack, expr: &Expr) {
             return;
         }
 
+        let ScopeKind::Function(function_def) = &checker.current_scope().kind else {
+            error!("Expected current scope kind to be FunctionDef");
+            return;
+        };
+        if function_def.globals.contains_key(id.as_str()) {
+            return;
+        }
+
         checker.add_check(Check::new(
             CheckKind::UnnecessaryAssign,
             Range::from_located(expr),