Improve reference resolution for deferred-annotations-within-classes

crates/ruff/resources/test/fixtures/pyflakes/F401_12.py

@@ -0,0 +1,10 @@
+"""Test: module bindings are preferred over local bindings, for deferred annotations."""
+
+from __future__ import annotations
+
+import datetime
+from typing import Optional
+
+
+class Class:
+    datetime: Optional[datetime.datetime]

crates/ruff/resources/test/fixtures/pyflakes/F401_13.py

@@ -0,0 +1,12 @@
+"""Test: module bindings are preferred over local bindings, for deferred annotations."""
+
+from __future__ import annotations
+
+from typing import TypeAlias, List
+
+
+class Class:
+    List: TypeAlias = List
+
+    def bar(self) -> List:
+        pass

crates/ruff/resources/test/fixtures/pyflakes/F401_14.py

@@ -0,0 +1,8 @@
+"""Test: module bindings are preferred over local bindings, for deferred annotations."""
+
+import datetime
+from typing import Optional
+
+
+class Class:
+    datetime: "Optional[datetime.datetime]"

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};
@@ -4732,40 +4732,20 @@ impl<'a> Checker<'a> {
         let Expr::Name(ast::ExprName { id, .. } )= expr else {
             return;
         };
-        let id = id.as_str();
-
-        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;
-                }
-
+        match self.ctx.resolve_reference(id, expr.range()) {
+            ResolvedReference::Resolved(scope_id, binding_id) => {
                 // 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:
-                //   import pyarrow as pa
-                //   import pyarrow.csv
-                //   print(pa.csv.read_csv("test.csv"))
-                match &self.ctx.bindings[*index].kind {
+                // 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[binding_id].kind {
                     BindingKind::Importation(Importation { name, full_name })
                     | BindingKind::SubmoduleImportation(SubmoduleImportation { name, full_name }) =>
                     {
@@ -4776,8 +4756,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) {
-                                self.ctx.bindings[*index].mark_used(
+                            if let Some(binding_id) = self.ctx.scopes[scope_id].get(full_name) {
+                                let context = self.ctx.execution_context();
+                                self.ctx.bindings[*binding_id].mark_used(
                                     self.ctx.scope_id,
                                     expr.range(),
                                     context,
@@ -4793,8 +4774,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()) {
-                                self.ctx.bindings[*index].mark_used(
+                            if let Some(binding_id) =
+                                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,
@@ -4804,72 +4788,70 @@ impl<'a> Checker<'a> {
                     }
                     _ => {}
                 }
-
-                return;
             }
-
-            first_iter = false;
-            import_starred = import_starred || scope.uses_star_imports();
-        }
-
-        if import_starred {
-            // F405
-            if self
-                .settings
-                .rules
-                .enabled(Rule::UndefinedLocalWithImportStarUsage)
-            {
-                let sources: Vec<String> = self
-                    .ctx
-                    .scopes
-                    .iter()
-                    .flat_map(Scope::star_imports)
-                    .map(|StarImportation { level, module }| {
-                        helpers::format_import_from(*level, *module)
-                    })
-                    .sorted()
-                    .dedup()
-                    .collect();
-                self.diagnostics.push(Diagnostic::new(
-                    pyflakes::rules::UndefinedLocalWithImportStarUsage {
-                        name: id.to_string(),
-                        sources,
-                    },
-                    expr.range(),
-                ));
+            ResolvedReference::ImplicitGlobal => {
+                // Nothing to do.
             }
-            return;
-        }
-
-        if self.settings.rules.enabled(Rule::UndefinedName) {
-            // Allow __path__.
-            if self.path.ends_with("__init__.py") && id == "__path__" {
-                return;
+            ResolvedReference::StarImport => {
+                // F405
+                if self
+                    .settings
+                    .rules
+                    .enabled(Rule::UndefinedLocalWithImportStarUsage)
+                {
+                    let sources: Vec<String> = self
+                        .ctx
+                        .scopes
+                        .iter()
+                        .flat_map(Scope::star_imports)
+                        .map(|StarImportation { level, module }| {
+                            helpers::format_import_from(*level, *module)
+                        })
+                        .sorted()
+                        .dedup()
+                        .collect();
+                    self.diagnostics.push(Diagnostic::new(
+                        pyflakes::rules::UndefinedLocalWithImportStarUsage {
+                            name: id.to_string(),
+                            sources,
+                        },
+                        expr.range(),
+                    ));
+                }
             }
+            ResolvedReference::NotFound => {
+                // F821
+                if self.settings.rules.enabled(Rule::UndefinedName) {
+                    // Allow __path__.
+                    if self.path.ends_with("__init__.py") && id == "__path__" {
+                        return;
+                    }
 
-            // Allow "__module__" and "__qualname__" in class scopes.
-            if (id == "__module__" || id == "__qualname__")
-                && matches!(self.ctx.scope().kind, ScopeKind::Class(..))
-            {
-                return;
-            }
+                    // Allow "__module__" and "__qualname__" in class scopes.
+                    if (id == "__module__" || id == "__qualname__")
+                        && matches!(self.ctx.scope().kind, ScopeKind::Class(..))
+                    {
+                        return;
+                    }
 
-            // Avoid flagging if NameError is handled.
-            if self
-                .ctx
-                .handled_exceptions
-                .iter()
-                .any(|handler_names| handler_names.contains(Exceptions::NAME_ERROR))
-            {
-                return;
-            }
+                    // Avoid flagging if `NameError` is handled.
+                    if self
+                        .ctx
+                        .handled_exceptions
+                        .iter()
+                        .any(|handler_names| handler_names.contains(Exceptions::NAME_ERROR))
+                    {
+                        return;
+                    }
 
-            self.diagnostics.push(Diagnostic::new(
-                pyflakes::rules::UndefinedName {
-                    name: id.to_string(),
-                },
-                expr.range(),
-            ));
+                    self.diagnostics.push(Diagnostic::new(
+                        pyflakes::rules::UndefinedName {
+                            name: id.to_string(),
+                        },
+                        expr.range(),
+                    ));
+                }
+            }
         }
     }
 

crates/ruff/src/rules/pyflakes/mod.rs

@@ -36,6 +36,9 @@ mod tests {
     #[test_case(Rule::UnusedImport, Path::new("F401_9.py"); "F401_9")]
     #[test_case(Rule::UnusedImport, Path::new("F401_10.py"); "F401_10")]
     #[test_case(Rule::UnusedImport, Path::new("F401_11.py"); "F401_11")]
+    #[test_case(Rule::UnusedImport, Path::new("F401_12.py"); "F401_12")]
+    #[test_case(Rule::UnusedImport, Path::new("F401_13.py"); "F401_13")]
+    #[test_case(Rule::UnusedImport, Path::new("F401_14.py"); "F401_14")]
     #[test_case(Rule::ImportShadowedByLoopVar, Path::new("F402.py"); "F402")]
     #[test_case(Rule::UndefinedLocalWithImportStar, Path::new("F403.py"); "F403")]
     #[test_case(Rule::LateFutureImport, Path::new("F404.py"); "F404")]
@@ -3175,8 +3178,8 @@ mod tests {
     fn unassigned_annotation_is_undefined() {
         flakes(
             r#"
-        name: str
-        print(name)
+name: str
+print(name)
         "#,
             &[Rule::UndefinedName],
         );

crates/ruff/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F401_F401_12.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff/src/rules/pyflakes/mod.rs
+---
+

crates/ruff/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F401_F401_13.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff/src/rules/pyflakes/mod.rs
+---
+

crates/ruff/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F401_F401_14.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff/src/rules/pyflakes/mod.rs
+---
+

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,66 @@ impl<'a> Context<'a> {
             .map_or(false, |binding| binding.kind.is_builtin())
     }
 
+    /// Resolve a read reference to the given symbol.
+    pub fn resolve_reference(&mut self, symbol: &str, range: TextRange) -> ResolvedReference {
+        // PEP 563 indicates that if a forward reference can be resolved in the module scope, we
+        // should prefer it over local resolutions.
+        if self.in_deferred_type_definition() {
+            let scope_id = ScopeId::global();
+            if let Some(binding_id) = self.scopes[scope_id].get(symbol) {
+                let context = self.execution_context();
+                self.bindings[*binding_id].mark_used(scope_id, range, context);
+                return ResolvedReference::Resolved(scope_id, *binding_id);
+            }
+        }
+
+        let mut first_iter = true;
+        let mut import_starred = false;
+        for scope_id in self.scopes.ancestor_ids(self.scope_id) {
+            let scope = &self.scopes[scope_id];
+            if scope.kind.is_class() {
+                if symbol == "__class__" {
+                    return ResolvedReference::ImplicitGlobal;
+                } else if !first_iter {
+                    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);
+            }
+
+            first_iter = false;
+            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.
     ///
@@ -701,10 +762,24 @@ impl ContextFlags {
 }
 
 /// A snapshot of the [`Context`] at a given point in the AST traversal.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct Snapshot {
     scope_id: ScopeId,
     stmt_id: Option<NodeId>,
     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,
+}