Move reference-resolution into Context (#4510)

crates/ruff/src/checkers/ast/mod.rs

@@ -27,7 +27,7 @@ use ruff_python_semantic::binding::{
     Binding, BindingId, BindingKind, Exceptions, ExecutionContext, Export, FromImportation,
     Importation, StarImportation, SubmoduleImportation,
 };
-use ruff_python_semantic::context::{Context, ContextFlags};
+use ruff_python_semantic::context::{Context, ContextFlags, ResolvedReference};
 use ruff_python_semantic::definition::{ContextualizedDefinition, Module, ModuleKind};
 use ruff_python_semantic::node::NodeId;
 use ruff_python_semantic::scope::{ClassDef, FunctionDef, Lambda, Scope, ScopeId, ScopeKind};
@@ -4746,40 +4746,20 @@ impl<'a> Checker<'a> {
         let Expr::Name(ast::ExprName { id, .. } )= expr else {
             return;
         };
-        let id = id.as_str();
+        match self.ctx.resolve_reference(id, expr.range()) {
-
+            ResolvedReference::Resolved(scope_id, binding_id) => {
-        let mut first_iter = true;
-        let mut import_starred = false;
-
-        for scope in self.ctx.scopes.ancestors(self.ctx.scope_id) {
-            if scope.kind.is_class() {
-                if id == "__class__" {
-                    return;
-                } else if !first_iter {
-                    continue;
-                }
-            }
-
-            if let Some(index) = scope.get(id) {
-                // Mark the binding as used.
-                let context = self.ctx.execution_context();
-                self.ctx.bindings[*index].mark_used(self.ctx.scope_id, expr.range(), context);
-
-                if !self.ctx.in_deferred_type_definition()
-                    && self.ctx.bindings[*index].kind.is_annotation()
-                {
-                    continue;
-                }
-
                 // If the name of the sub-importation is the same as an alias of another
                 // importation and the alias is used, that sub-importation should be
                 // marked as used too.
                 //
-                // This handles code like:
+                // For example, mark `pa` as used in:
+                //
+                // ```python
                 // import pyarrow as pa
                 // import pyarrow.csv
                 // print(pa.csv.read_csv("test.csv"))
-                match &self.ctx.bindings[*index].kind {
+                // ```
+                match &self.ctx.bindings[binding_id].kind {
                     BindingKind::Importation(Importation { name, full_name })
                     | BindingKind::SubmoduleImportation(SubmoduleImportation { name, full_name }) =>
                     {
@@ -4790,8 +4770,9 @@ impl<'a> Checker<'a> {
                             .unwrap_or_default();
                         if has_alias {
                             // Mark the sub-importation as used.
-                            if let Some(index) = scope.get(full_name) {
+                            if let Some(binding_id) = self.ctx.scopes[scope_id].get(full_name) {
-                                self.ctx.bindings[*index].mark_used(
+                                let context = self.ctx.execution_context();
+                                self.ctx.bindings[*binding_id].mark_used(
                                     self.ctx.scope_id,
                                     expr.range(),
                                     context,
@@ -4807,8 +4788,11 @@ impl<'a> Checker<'a> {
                             .unwrap_or_default();
                         if has_alias {
                             // Mark the sub-importation as used.
-                            if let Some(index) = scope.get(full_name.as_str()) {
+                            if let Some(binding_id) =
-                                self.ctx.bindings[*index].mark_used(
+                                self.ctx.scopes[scope_id].get(full_name.as_str())
+                            {
+                                let context = self.ctx.execution_context();
+                                self.ctx.bindings[*binding_id].mark_used(
                                     self.ctx.scope_id,
                                     expr.range(),
                                     context,
@@ -4818,15 +4802,11 @@ impl<'a> Checker<'a> {
                     }
                     _ => {}
                 }
-
-                return;
             }
-
+            ResolvedReference::ImplicitGlobal => {
-            first_iter = false;
+                // Nothing to do.
-            import_starred = import_starred || scope.uses_star_imports();
             }
-
+            ResolvedReference::StarImport => {
-        if import_starred {
                 // F405
                 if self
                     .settings
@@ -4852,9 +4832,9 @@ impl<'a> Checker<'a> {
                         expr.range(),
                     ));
                 }
-            return;
             }
-
+            ResolvedReference::NotFound => {
+                // F821
                 if self.settings.rules.enabled(Rule::UndefinedName) {
                     // Allow __path__.
                     if self.path.ends_with("__init__.py") && id == "__path__" {
@@ -4868,7 +4848,7 @@ impl<'a> Checker<'a> {
                         return;
                     }
 
-            // Avoid flagging if NameError is handled.
+                    // Avoid flagging if `NameError` is handled.
                     if self
                         .ctx
                         .handled_exceptions
@@ -4886,6 +4866,8 @@ impl<'a> Checker<'a> {
                     ));
                 }
             }
+        }
+    }
 
     fn handle_node_store(&mut self, id: &'a str, expr: &Expr) {
         let parent = self.ctx.stmt();

crates/ruff_python_semantic/src/context.rs

@@ -3,6 +3,7 @@ use std::path::Path;
 
 use bitflags::bitflags;
 use nohash_hasher::{BuildNoHashHasher, IntMap};
+use ruff_text_size::TextRange;
 use rustpython_parser::ast::{Expr, Stmt};
 use smallvec::smallvec;
 
@@ -113,6 +114,54 @@ impl<'a> Context<'a> {
             .map_or(false, |binding| binding.kind.is_builtin())
     }
 
+    /// Resolve a reference to the given symbol.
+    pub fn resolve_reference(&mut self, symbol: &str, range: TextRange) -> ResolvedReference {
+        let mut import_starred = false;
+        for (index, scope_id) in self.scopes.ancestor_ids(self.scope_id).enumerate() {
+            let scope = &self.scopes[scope_id];
+            if scope.kind.is_class() {
+                if symbol == "__class__" {
+                    return ResolvedReference::ImplicitGlobal;
+                }
+                if index > 0 {
+                    continue;
+                }
+            }
+
+            if let Some(binding_id) = scope.get(symbol) {
+                // Mark the binding as used.
+                let context = self.execution_context();
+                self.bindings[*binding_id].mark_used(self.scope_id, range, context);
+
+                // But if it's a type annotation, don't treat it as resolved, unless we're in a
+                // forward reference. For example, given:
+                //
+                // ```python
+                // name: str
+                // print(name)
+                // ```
+                //
+                // The `name` in `print(name)` should be treated as unresolved, but the `name` in
+                // `name: str` should be treated as used.
+                if !self.in_deferred_type_definition()
+                    && self.bindings[*binding_id].kind.is_annotation()
+                {
+                    continue;
+                }
+
+                return ResolvedReference::Resolved(scope_id, *binding_id);
+            }
+
+            import_starred = import_starred || scope.uses_star_imports();
+        }
+
+        if import_starred {
+            ResolvedReference::StarImport
+        } else {
+            ResolvedReference::NotFound
+        }
+    }
+
     /// Resolves the [`Expr`] to a fully-qualified symbol-name, if `value` resolves to an imported
     /// or builtin symbol.
     ///
@@ -708,3 +757,17 @@ pub struct Snapshot {
     definition_id: DefinitionId,
     flags: ContextFlags,
 }
+
+#[derive(Debug)]
+pub enum ResolvedReference {
+    /// The reference is resolved to a specific binding.
+    Resolved(ScopeId, BindingId),
+    /// The reference is resolved to a context-specific, implicit global (e.g., `__class__` within
+    /// a class scope).
+    ImplicitGlobal,
+    /// The reference is unresolved, but at least one of the containing scopes contains a star
+    /// import.
+    StarImport,
+    /// The reference is definitively unresolved.
+    NotFound,
+}