From 66ae4db6cdca56cc2c2fb8e157473341ec6aa281 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Thu, 17 Nov 2022 17:19:01 -0500
Subject: [PATCH] Ignore globals when checking local variable names (#800)

---
 resources/test/fixtures/N806.py               |  4 ++
 src/ast/types.rs                              |  2 +
 src/check_ast.rs                              | 58 +++++++++----------
 src/pyflakes/checks.rs                        |  5 +-
 .../ruff__linter__tests__N806_N806.py.snap    |  8 +--
 5 files changed, 42 insertions(+), 35 deletions(-)

diff --git a/resources/test/fixtures/N806.py b/resources/test/fixtures/N806.py
index effd9d4cd2..861bce4baa 100644
--- a/resources/test/fixtures/N806.py
+++ b/resources/test/fixtures/N806.py
@@ -1,8 +1,12 @@
 import collections
 from collections import namedtuple
 
+GLOBAL: str = "foo"
+
 
 def f():
+    global GLOBAL
+    GLOBAL = "bar"
     lower = 0
     Camel = 0
     CONSTANT = 0
diff --git a/src/ast/types.rs b/src/ast/types.rs
index ad0b3a7c37..be94b1ee05 100644
--- a/src/ast/types.rs
+++ b/src/ast/types.rs
@@ -79,8 +79,10 @@ pub enum BindingKind {
     Annotation,
     Argument,
     Assignment,
+    // TODO(charlie): This seems to be a catch-all.
     Binding,
     LoopVar,
+    Global,
     Builtin,
     ClassDefinition,
     Definition,
diff --git a/src/check_ast.rs b/src/check_ast.rs
index f1ad631658..2b96911c30 100644
--- a/src/check_ast.rs
+++ b/src/check_ast.rs
@@ -19,8 +19,7 @@ use crate::ast::helpers::{
 use crate::ast::operations::extract_all_names;
 use crate::ast::relocate::relocate_expr;
 use crate::ast::types::{
-    Binding, BindingContext, BindingKind, ClassScope, FunctionScope, ImportKind, Range, Scope,
-    ScopeKind,
+    Binding, BindingContext, BindingKind, ClassScope, ImportKind, Range, Scope, ScopeKind,
 };
 use crate::ast::visitor::{walk_excepthandler, Visitor};
 use crate::ast::{helpers, operations, visitor};
@@ -216,21 +215,18 @@ where
         match &stmt.node {
             StmtKind::Global { names } | StmtKind::Nonlocal { names } => {
                 let global_scope_id = self.scopes[GLOBAL_SCOPE_INDEX].id;
-
-                let current_scope = self.current_scope();
-                let current_scope_id = current_scope.id;
-                if current_scope_id != global_scope_id {
+                let scope =
+                    &mut self.scopes[*(self.scope_stack.last().expect("No current scope found."))];
+                if scope.id != global_scope_id {
                     for name in names {
-                        for scope in self.scopes.iter_mut().skip(GLOBAL_SCOPE_INDEX + 1) {
-                            scope.values.insert(
-                                name,
-                                Binding {
-                                    kind: BindingKind::Assignment,
-                                    used: Some((global_scope_id, Range::from_located(stmt))),
-                                    range: Range::from_located(stmt),
-                                },
-                            );
-                        }
+                        scope.values.insert(
+                            name,
+                            Binding {
+                                kind: BindingKind::Global,
+                                used: Some((global_scope_id, Range::from_located(stmt))),
+                                range: Range::from_located(stmt),
+                            },
+                        );
                     }
                 }
 
@@ -732,10 +728,8 @@ where
                             ));
                         }
 
-                        let scope = &mut self.scopes[*(self
-                            .scope_stack
-                            .last_mut()
-                            .expect("No current scope found."))];
+                        let scope = &mut self.scopes
+                            [*(self.scope_stack.last().expect("No current scope found."))];
                         scope.import_starred = true;
                     } else {
                         if let Some(asname) = &alias.node.asname {
@@ -1452,15 +1446,10 @@ where
 
                 if let ExprKind::Name { id, ctx } = &func.node {
                     if id == "locals" && matches!(ctx, ExprContext::Load) {
-                        let scope = &mut self.scopes[*(self
-                            .scope_stack
-                            .last_mut()
-                            .expect("No current scope found."))];
-                        if matches!(
-                            scope.kind,
-                            ScopeKind::Function(FunctionScope { uses_locals: false })
-                        ) {
-                            scope.kind = ScopeKind::Function(FunctionScope { uses_locals: true });
+                        let scope = &mut self.scopes
+                            [*(self.scope_stack.last().expect("No current scope found."))];
+                        if let ScopeKind::Function(inner) = &mut scope.kind {
+                            inner.uses_locals = true;
                         }
                     }
                 }
@@ -2243,7 +2232,16 @@ impl<'a> Checker<'a> {
 
         if self.settings.enabled.contains(&CheckCode::N806) {
             if matches!(self.current_scope().kind, ScopeKind::Function(..)) {
-                pep8_naming::plugins::non_lowercase_variable_in_function(self, expr, parent, id)
+                // Ignore globals.
+                if !self
+                    .current_scope()
+                    .values
+                    .get(id)
+                    .map(|binding| matches!(binding.kind, BindingKind::Global))
+                    .unwrap_or(false)
+                {
+                    pep8_naming::plugins::non_lowercase_variable_in_function(self, expr, parent, id)
+                }
             }
         }
 
diff --git a/src/pyflakes/checks.rs b/src/pyflakes/checks.rs
index 7e116121d3..cda84d5057 100644
--- a/src/pyflakes/checks.rs
+++ b/src/pyflakes/checks.rs
@@ -55,7 +55,10 @@ pub fn unused_variables(scope: &Scope, dummy_variable_rgx: &Regex) -> Vec<Check>
 
     if matches!(
         scope.kind,
-        ScopeKind::Function(FunctionScope { uses_locals: true })
+        ScopeKind::Function(FunctionScope {
+            uses_locals: true,
+            ..
+        })
     ) {
         return checks;
     }
diff --git a/src/snapshots/ruff__linter__tests__N806_N806.py.snap b/src/snapshots/ruff__linter__tests__N806_N806.py.snap
index aa4b4f8a14..e6e8725df4 100644
--- a/src/snapshots/ruff__linter__tests__N806_N806.py.snap
+++ b/src/snapshots/ruff__linter__tests__N806_N806.py.snap
@@ -5,19 +5,19 @@ expression: checks
 - kind:
     NonLowercaseVariableInFunction: Camel
   location:
-    row: 7
+    row: 11
     column: 4
   end_location:
-    row: 7
+    row: 11
     column: 9
   fix: ~
 - kind:
     NonLowercaseVariableInFunction: CONSTANT
   location:
-    row: 8
+    row: 12
     column: 4
   end_location:
-    row: 8
+    row: 12
     column: 12
   fix: ~