Improve ruff_python_semantic::all::extract_all_names() (#11335)

Author: AlexWaygood

crates/ruff_linter/resources/test/fixtures/pyflakes/F822_1b.py

@@ -0,0 +1,4 @@
+import builtins
+a = 1
+
+__all__ = builtins.list(["a", "b"])

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

@@ -49,7 +49,7 @@ use ruff_python_ast::{helpers, str, visitor, PySourceType};
 use ruff_python_codegen::{Generator, Stylist};
 use ruff_python_index::Indexer;
 use ruff_python_parser::typing::{parse_type_annotation, AnnotationKind};
-use ruff_python_semantic::all::{extract_all_names, DunderAllDefinition, DunderAllFlags};
+use ruff_python_semantic::all::{DunderAllDefinition, DunderAllFlags};
 use ruff_python_semantic::analyze::{imports, typing};
 use ruff_python_semantic::{
     BindingFlags, BindingId, BindingKind, Exceptions, Export, FromImport, Globals, Import, Module,
@@ -1882,8 +1882,7 @@ impl<'a> Checker<'a> {
                 _ => false,
             }
         {
-            let (all_names, all_flags) =
-                extract_all_names(parent, |name| self.semantic.has_builtin_binding(name));
+            let (all_names, all_flags) = self.semantic.extract_dunder_all_names(parent);
 
             if all_flags.intersects(DunderAllFlags::INVALID_OBJECT) {
                 flags |= BindingFlags::INVALID_ALL_OBJECT;

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

@@ -161,6 +161,7 @@ mod tests {
     #[test_case(Rule::UndefinedExport, Path::new("F822_0.py"))]
     #[test_case(Rule::UndefinedExport, Path::new("F822_0.pyi"))]
     #[test_case(Rule::UndefinedExport, Path::new("F822_1.py"))]
+    #[test_case(Rule::UndefinedExport, Path::new("F822_1b.py"))]
     #[test_case(Rule::UndefinedExport, Path::new("F822_2.py"))]
     #[test_case(Rule::UndefinedExport, Path::new("F822_3.py"))]
     #[test_case(Rule::UndefinedLocal, Path::new("F823.py"))]

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F822_F822_1b.py.snap

@@ -0,0 +1,10 @@
+---
+source: crates/ruff_linter/src/rules/pyflakes/mod.rs
+---
+F822_1b.py:4:31: F822 Undefined name `b` in `__all__`
+  |
+2 | a = 1
+3 | 
+4 | __all__ = builtins.list(["a", "b"])
+  |                               ^^^ F822
+  |

crates/ruff_python_semantic/src/definition.rs

@@ -5,14 +5,14 @@ use std::fmt::Debug;
 use std::ops::Deref;
 use std::path::Path;
 
-use crate::all::DunderAllName;
 use ruff_index::{newtype_index, IndexSlice, IndexVec};
 use ruff_python_ast::{self as ast, Stmt};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::analyze::visibility::{
     class_visibility, function_visibility, method_visibility, module_visibility, Visibility,
 };
+use crate::model::all::DunderAllName;
 
 /// Id uniquely identifying a definition in a program.
 #[newtype_index]

crates/ruff_python_semantic/src/lib.rs

@@ -1,4 +1,3 @@
-pub mod all;
 pub mod analyze;
 mod binding;
 mod branches;

crates/ruff_python_semantic/src/model.rs

@@ -1,3 +1,5 @@
+pub mod all;
+
 use std::path::Path;
 
 use bitflags::bitflags;

crates/ruff_python_semantic/src/all.rs renamed to crates/ruff_python_semantic/src/model/all.rs

@@ -1,8 +1,12 @@
+//! Utilities for semantic analysis of `__all__` definitions
+
 use bitflags::bitflags;
 
 use ruff_python_ast::{self as ast, helpers::map_subscript, Expr, Stmt};
 use ruff_text_size::{Ranged, TextRange};
 
+use crate::SemanticModel;
+
 bitflags! {
     #[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
     pub struct DunderAllFlags: u8 {
@@ -66,34 +70,79 @@ impl Ranged for DunderAllDefinition<'_> {
     }
 }
 
-/// Extract the names bound to a given __all__ assignment.
-///
-/// Accepts a closure that determines whether a given name (e.g., `"list"`) is a Python builtin.
-pub fn extract_all_names<F>(stmt: &Stmt, is_builtin: F) -> (Vec<DunderAllName>, DunderAllFlags)
-where
-    F: Fn(&str) -> bool,
-{
-    fn add_to_names<'a>(
-        elts: &'a [Expr],
-        names: &mut Vec<DunderAllName<'a>>,
-        flags: &mut DunderAllFlags,
-    ) {
-        for elt in elts {
-            if let Expr::StringLiteral(ast::ExprStringLiteral { value, range }) = elt {
-                names.push(DunderAllName {
-                    name: value.to_str(),
-                    range: *range,
-                });
+impl<'a> SemanticModel<'a> {
+    /// Extract the names bound to a given __all__ assignment.
+    pub fn extract_dunder_all_names<'expr>(
+        &self,
+        stmt: &'expr Stmt,
+    ) -> (Vec<DunderAllName<'expr>>, DunderAllFlags) {
+        fn add_to_names<'expr>(
+            elts: &'expr [Expr],
+            names: &mut Vec<DunderAllName<'expr>>,
+            flags: &mut DunderAllFlags,
+        ) {
+            for elt in elts {
+                if let Expr::StringLiteral(ast::ExprStringLiteral { value, range }) = elt {
+                    names.push(DunderAllName {
+                        name: value.to_str(),
+                        range: *range,
+                    });
+                } else {
+                    *flags |= DunderAllFlags::INVALID_OBJECT;
+                }
+            }
+        }
+
+        let mut names: Vec<DunderAllName> = vec![];
+        let mut flags = DunderAllFlags::empty();
+
+        if let Some(value) = match stmt {
+            Stmt::Assign(ast::StmtAssign { value, .. }) => Some(value),
+            Stmt::AnnAssign(ast::StmtAnnAssign { value, .. }) => value.as_ref(),
+            Stmt::AugAssign(ast::StmtAugAssign { value, .. }) => Some(value),
+            _ => None,
+        } {
+            if let Expr::BinOp(ast::ExprBinOp { left, right, .. }) = value.as_ref() {
+                let mut current_left = left;
+                let mut current_right = right;
+                loop {
+                    // Process the right side, which should be a "real" value.
+                    let (elts, new_flags) = self.extract_dunder_all_elts(current_right);
+                    flags |= new_flags;
+                    if let Some(elts) = elts {
+                        add_to_names(elts, &mut names, &mut flags);
+                    }
+
+                    // Process the left side, which can be a "real" value or the "rest" of the
+                    // binary operation.
+                    if let Expr::BinOp(ast::ExprBinOp { left, right, .. }) = current_left.as_ref() {
+                        current_left = left;
+                        current_right = right;
+                    } else {
+                        let (elts, new_flags) = self.extract_dunder_all_elts(current_left);
+                        flags |= new_flags;
+                        if let Some(elts) = elts {
+                            add_to_names(elts, &mut names, &mut flags);
+                        }
+                        break;
+                    }
+                }
             } else {
-                *flags |= DunderAllFlags::INVALID_OBJECT;
+                let (elts, new_flags) = self.extract_dunder_all_elts(value);
+                flags |= new_flags;
+                if let Some(elts) = elts {
+                    add_to_names(elts, &mut names, &mut flags);
+                }
             }
         }
+
+        (names, flags)
     }
 
-    fn extract_elts<F>(expr: &Expr, is_builtin: F) -> (Option<&[Expr]>, DunderAllFlags)
-    where
-        F: Fn(&str) -> bool,
-    {
+    fn extract_dunder_all_elts<'expr>(
+        &self,
+        expr: &'expr Expr,
+    ) -> (Option<&'expr [Expr]>, DunderAllFlags) {
         match expr {
             Expr::List(ast::ExprList { elts, .. }) => {
                 return (Some(elts), DunderAllFlags::empty());
@@ -122,80 +171,35 @@ where
             }) => {
                 // Allow `tuple()`, `list()`, and their generic forms, like `list[int]()`.
                 if arguments.keywords.is_empty() && arguments.args.len() <= 1 {
-                    if let Expr::Name(ast::ExprName { id, .. }) = map_subscript(func) {
-                        let id = id.as_str();
-                        if matches!(id, "tuple" | "list") && is_builtin(id) {
-                            let [arg] = arguments.args.as_ref() else {
+                    if self
+                        .resolve_builtin_symbol(map_subscript(func))
+                        .is_some_and(|symbol| matches!(symbol, "tuple" | "list"))
+                    {
+                        let [arg] = arguments.args.as_ref() else {
+                            return (None, DunderAllFlags::empty());
+                        };
+                        match arg {
+                            Expr::List(ast::ExprList { elts, .. })
+                            | Expr::Set(ast::ExprSet { elts, .. })
+                            | Expr::Tuple(ast::ExprTuple { elts, .. }) => {
+                                return (Some(elts), DunderAllFlags::empty());
+                            }
+                            _ => {
+                                // We can't analyze other expressions, but they must be
+                                // valid, since the `list` or `tuple` call will ultimately
+                                // evaluate to a list or tuple.
                                 return (None, DunderAllFlags::empty());
-                            };
-                            match arg {
-                                Expr::List(ast::ExprList { elts, .. })
-                                | Expr::Set(ast::ExprSet { elts, .. })
-                                | Expr::Tuple(ast::ExprTuple { elts, .. }) => {
-                                    return (Some(elts), DunderAllFlags::empty());
-                                }
-                                _ => {
-                                    // We can't analyze other expressions, but they must be
-                                    // valid, since the `list` or `tuple` call will ultimately
-                                    // evaluate to a list or tuple.
-                                    return (None, DunderAllFlags::empty());
-                                }
                             }
                         }
                     }
                 }
             }
             Expr::Named(ast::ExprNamed { value, .. }) => {
                 // Allow, e.g., `__all__ += (value := ["A", "B"])`.
-                return extract_elts(value, is_builtin);
+                return self.extract_dunder_all_elts(value);
             }
             _ => {}
         }
         (None, DunderAllFlags::INVALID_FORMAT)
     }
-
-    let mut names: Vec<DunderAllName> = vec![];
-    let mut flags = DunderAllFlags::empty();
-
-    if let Some(value) = match stmt {
-        Stmt::Assign(ast::StmtAssign { value, .. }) => Some(value),
-        Stmt::AnnAssign(ast::StmtAnnAssign { value, .. }) => value.as_ref(),
-        Stmt::AugAssign(ast::StmtAugAssign { value, .. }) => Some(value),
-        _ => None,
-    } {
-        if let Expr::BinOp(ast::ExprBinOp { left, right, .. }) = value.as_ref() {
-            let mut current_left = left;
-            let mut current_right = right;
-            loop {
-                // Process the right side, which should be a "real" value.
-                let (elts, new_flags) = extract_elts(current_right, |expr| is_builtin(expr));
-                flags |= new_flags;
-                if let Some(elts) = elts {
-                    add_to_names(elts, &mut names, &mut flags);
-                }
-
-                // Process the left side, which can be a "real" value or the "rest" of the
-                // binary operation.
-                if let Expr::BinOp(ast::ExprBinOp { left, right, .. }) = current_left.as_ref() {
-                    current_left = left;
-                    current_right = right;
-                } else {
-                    let (elts, new_flags) = extract_elts(current_left, |expr| is_builtin(expr));
-                    flags |= new_flags;
-                    if let Some(elts) = elts {
-                        add_to_names(elts, &mut names, &mut flags);
-                    }
-                    break;
-                }
-            }
-        } else {
-            let (elts, new_flags) = extract_elts(value, |expr| is_builtin(expr));
-            flags |= new_flags;
-            if let Some(elts) = elts {
-                add_to_names(elts, &mut names, &mut flags);
-            }
-        }
-    }
-
-    (names, flags)
 }