From 780d153ae80139845597ec0b853ad682d0c2281a Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Tue, 6 Jun 2023 21:22:21 -0400
Subject: [PATCH] Replace one-off locals property with `ScopeFlags` (#4912)

---
 crates/ruff/src/checkers/ast/mod.rs           |  2 +-
 .../rules/pyflakes/rules/unused_variable.rs   |  4 +--
 crates/ruff_python_semantic/src/scope.rs      | 28 ++++++++++++++++---
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
index 25b2e30af0..03165631dc 100644
--- a/crates/ruff/src/checkers/ast/mod.rs
+++ b/crates/ruff/src/checkers/ast/mod.rs
@@ -2815,7 +2815,7 @@ where
                 if let Expr::Name(ast::ExprName { id, ctx, range: _ }) = func.as_ref() {
                     if id == "locals" && matches!(ctx, ExprContext::Load) {
                         let scope = self.semantic_model.scope_mut();
-                        scope.uses_locals = true;
+                        scope.set_uses_locals();
                     }
                 }
 
diff --git a/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs b/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
index 67b201f83b..fb6932ca3f 100644
--- a/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
+++ b/crates/ruff/src/rules/pyflakes/rules/unused_variable.rs
@@ -7,7 +7,7 @@ use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::contains_effect;
 use ruff_python_ast::source_code::Locator;
-use ruff_python_semantic::scope::{ScopeId, ScopeKind};
+use ruff_python_semantic::scope::ScopeId;
 
 use crate::autofix::edits::delete_stmt;
 use crate::checkers::ast::Checker;
@@ -286,7 +286,7 @@ fn remove_unused_variable(
 /// F841
 pub(crate) fn unused_variable(checker: &mut Checker, scope: ScopeId) {
     let scope = &checker.semantic_model().scopes[scope];
-    if scope.uses_locals && matches!(scope.kind, ScopeKind::Function(..)) {
+    if scope.uses_locals() && scope.kind.is_any_function() {
         return;
     }
 
diff --git a/crates/ruff_python_semantic/src/scope.rs b/crates/ruff_python_semantic/src/scope.rs
index be3edffe05..e56d4dafa8 100644
--- a/crates/ruff_python_semantic/src/scope.rs
+++ b/crates/ruff_python_semantic/src/scope.rs
@@ -1,3 +1,4 @@
+use bitflags::bitflags;
 use nohash_hasher::{BuildNoHashHasher, IntMap};
 use std::collections::HashMap;
 use std::ops::{Deref, DerefMut};
@@ -14,8 +15,6 @@ use crate::globals::GlobalsId;
 pub struct Scope<'a> {
     pub kind: ScopeKind<'a>,
     pub parent: Option<ScopeId>,
-    /// Whether this scope uses the `locals()` builtin.
-    pub uses_locals: bool,
     /// A list of star imports in this scope. These represent _module_ imports (e.g., `sys` in
     /// `from sys import *`), rather than individual bindings (e.g., individual members in `sys`).
     star_imports: Vec<StarImportation<'a>>,
@@ -27,6 +26,8 @@ pub struct Scope<'a> {
     deleted_symbols: Vec<&'a str>,
     /// Index into the globals arena, if the scope contains any globally-declared symbols.
     globals_id: Option<GlobalsId>,
+    /// Flags for the [`Scope`].
+    flags: ScopeFlags,
 }
 
 impl<'a> Scope<'a> {
@@ -34,12 +35,12 @@ impl<'a> Scope<'a> {
         Scope {
             kind: ScopeKind::Module,
             parent: None,
-            uses_locals: false,
             star_imports: Vec::default(),
             bindings: FxHashMap::default(),
             shadowed_bindings: IntMap::default(),
             deleted_symbols: Vec::default(),
             globals_id: None,
+            flags: ScopeFlags::empty(),
         }
     }
 
@@ -47,12 +48,12 @@ impl<'a> Scope<'a> {
         Scope {
             kind,
             parent: Some(parent),
-            uses_locals: false,
             star_imports: Vec::default(),
             bindings: FxHashMap::default(),
             shadowed_bindings: IntMap::default(),
             deleted_symbols: Vec::default(),
             globals_id: None,
+            flags: ScopeFlags::empty(),
         }
     }
 
@@ -132,6 +133,25 @@ impl<'a> Scope<'a> {
     pub fn globals_id(&self) -> Option<GlobalsId> {
         self.globals_id
     }
+
+    /// Sets the [`ScopeFlags::USES_LOCALS`] flag.
+    pub fn set_uses_locals(&mut self) {
+        self.flags.insert(ScopeFlags::USES_LOCALS);
+    }
+
+    /// Returns `true` if this scope uses locals (e.g., `locals()`).
+    pub const fn uses_locals(&self) -> bool {
+        self.flags.contains(ScopeFlags::USES_LOCALS)
+    }
+}
+
+bitflags! {
+    /// Flags on a [`Scope`].
+    #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
+    pub struct ScopeFlags: u8 {
+        /// The scope uses locals (e.g., `locals()`).
+        const USES_LOCALS = 1 << 0;
+    }
 }
 
 #[derive(Debug, is_macro::Is)]