diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs index a4adc45a0b..fe20864cd0 100644 --- a/crates/ruff/src/checkers/ast/mod.rs +++ b/crates/ruff/src/checkers/ast/mod.rs @@ -13,9 +13,9 @@ use rustpython_parser::ast::{ }; use ruff_diagnostics::Diagnostic; +use ruff_python_ast::all::{extract_all_names, AllNamesFlags}; use ruff_python_ast::context::Context; use ruff_python_ast::helpers::{extract_handled_exceptions, to_module_path}; -use ruff_python_ast::operations::{extract_all_names, AllNamesFlags}; use ruff_python_ast::scope::{ Binding, BindingId, BindingKind, ClassDef, Exceptions, ExecutionContext, Export, FromImportation, FunctionDef, Importation, Lambda, Scope, ScopeId, ScopeKind, ScopeStack, diff --git a/crates/ruff_python_ast/src/all.rs b/crates/ruff_python_ast/src/all.rs new file mode 100644 index 0000000000..998ab314d8 --- /dev/null +++ b/crates/ruff_python_ast/src/all.rs @@ -0,0 +1,141 @@ +use bitflags::bitflags; +use rustpython_parser::ast::{Constant, Expr, ExprKind, Stmt, StmtKind}; + +use crate::context::Context; +use crate::scope::{BindingKind, Export, Scope}; + +bitflags! { + #[derive(Default)] + pub struct AllNamesFlags: u32 { + const INVALID_FORMAT = 0b0000_0001; + const INVALID_OBJECT = 0b0000_0010; + } +} + +/// Extract the names bound to a given __all__ assignment. +pub fn extract_all_names( + ctx: &Context, + stmt: &Stmt, + scope: &Scope, +) -> (Vec, AllNamesFlags) { + fn add_to_names(names: &mut Vec, elts: &[Expr], flags: &mut AllNamesFlags) { + for elt in elts { + if let ExprKind::Constant { + value: Constant::Str(value), + .. + } = &elt.node + { + names.push(value.to_string()); + } else { + *flags |= AllNamesFlags::INVALID_OBJECT; + } + } + } + + fn extract_elts<'a>( + ctx: &'a Context, + expr: &'a Expr, + ) -> (Option<&'a Vec>, AllNamesFlags) { + match &expr.node { + ExprKind::List { elts, .. } => { + return (Some(elts), AllNamesFlags::empty()); + } + ExprKind::Tuple { elts, .. } => { + return (Some(elts), AllNamesFlags::empty()); + } + ExprKind::ListComp { .. } => { + // Allow comprehensions, even though we can't statically analyze them. + return (None, AllNamesFlags::empty()); + } + ExprKind::Call { + func, + args, + keywords, + .. + } => { + // Allow `tuple()` and `list()` calls. + if keywords.is_empty() && args.len() <= 1 { + if ctx.resolve_call_path(func).map_or(false, |call_path| { + call_path.as_slice() == ["", "tuple"] + || call_path.as_slice() == ["", "list"] + }) { + if args.is_empty() { + return (None, AllNamesFlags::empty()); + } + match &args[0].node { + ExprKind::List { elts, .. } + | ExprKind::Set { elts, .. } + | ExprKind::Tuple { elts, .. } => { + return (Some(elts), AllNamesFlags::empty()); + } + ExprKind::ListComp { .. } + | ExprKind::SetComp { .. } + | ExprKind::GeneratorExp { .. } => { + // Allow comprehensions, even though we can't statically analyze + // them. + return (None, AllNamesFlags::empty()); + } + _ => {} + } + } + } + } + _ => {} + } + (None, AllNamesFlags::INVALID_FORMAT) + } + + let mut names: Vec = vec![]; + let mut flags = AllNamesFlags::empty(); + + // Grab the existing bound __all__ values. + if let StmtKind::AugAssign { .. } = &stmt.node { + if let Some(index) = scope.get("__all__") { + if let BindingKind::Export(Export { names: existing }) = &ctx.bindings[*index].kind { + names.extend_from_slice(existing); + } + } + } + + if let Some(value) = match &stmt.node { + StmtKind::Assign { value, .. } => Some(value), + StmtKind::AnnAssign { value, .. } => value.as_ref(), + StmtKind::AugAssign { value, .. } => Some(value), + _ => None, + } { + if let ExprKind::BinOp { left, right, .. } = &value.node { + 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(ctx, current_right); + flags |= new_flags; + if let Some(elts) = elts { + add_to_names(&mut names, elts, &mut flags); + } + + // Process the left side, which can be a "real" value or the "rest" of the + // binary operation. + if let ExprKind::BinOp { left, right, .. } = ¤t_left.node { + current_left = left; + current_right = right; + } else { + let (elts, new_flags) = extract_elts(ctx, current_left); + flags |= new_flags; + if let Some(elts) = elts { + add_to_names(&mut names, elts, &mut flags); + } + break; + } + } + } else { + let (elts, new_flags) = extract_elts(ctx, value); + flags |= new_flags; + if let Some(elts) = elts { + add_to_names(&mut names, elts, &mut flags); + } + } + } + + (names, flags) +} diff --git a/crates/ruff_python_ast/src/lib.rs b/crates/ruff_python_ast/src/lib.rs index f7f0223c29..eee19941b9 100644 --- a/crates/ruff_python_ast/src/lib.rs +++ b/crates/ruff_python_ast/src/lib.rs @@ -1,3 +1,4 @@ +pub mod all; pub mod branch_detection; pub mod cast; pub mod comparable; diff --git a/crates/ruff_python_ast/src/operations.rs b/crates/ruff_python_ast/src/operations.rs index d2be030d0b..c37b160630 100644 --- a/crates/ruff_python_ast/src/operations.rs +++ b/crates/ruff_python_ast/src/operations.rs @@ -1,150 +1,11 @@ -use bitflags::bitflags; use rustc_hash::FxHashMap; -use rustpython_parser::ast::{Cmpop, Constant, Expr, ExprKind, Located, Stmt, StmtKind}; +use rustpython_parser::ast::{Cmpop, Expr, ExprKind, Located, Stmt, StmtKind}; use rustpython_parser::{lexer, Mode, Tok}; -use crate::context::Context; use crate::helpers::any_over_expr; -use crate::scope::{BindingKind, Export, Scope}; use crate::visitor; use crate::visitor::Visitor; -bitflags! { - #[derive(Default)] - pub struct AllNamesFlags: u32 { - const INVALID_FORMAT = 0b0000_0001; - const INVALID_OBJECT = 0b0000_0010; - } -} - -/// Extract the names bound to a given __all__ assignment. -pub fn extract_all_names( - ctx: &Context, - stmt: &Stmt, - scope: &Scope, -) -> (Vec, AllNamesFlags) { - fn add_to_names(names: &mut Vec, elts: &[Expr], flags: &mut AllNamesFlags) { - for elt in elts { - if let ExprKind::Constant { - value: Constant::Str(value), - .. - } = &elt.node - { - names.push(value.to_string()); - } else { - *flags |= AllNamesFlags::INVALID_OBJECT; - } - } - } - - fn extract_elts<'a>( - ctx: &'a Context, - expr: &'a Expr, - ) -> (Option<&'a Vec>, AllNamesFlags) { - match &expr.node { - ExprKind::List { elts, .. } => { - return (Some(elts), AllNamesFlags::empty()); - } - ExprKind::Tuple { elts, .. } => { - return (Some(elts), AllNamesFlags::empty()); - } - ExprKind::ListComp { .. } => { - // Allow comprehensions, even though we can't statically analyze them. - return (None, AllNamesFlags::empty()); - } - ExprKind::Call { - func, - args, - keywords, - .. - } => { - // Allow `tuple()` and `list()` calls. - if keywords.is_empty() && args.len() <= 1 { - if ctx.resolve_call_path(func).map_or(false, |call_path| { - call_path.as_slice() == ["", "tuple"] - || call_path.as_slice() == ["", "list"] - }) { - if args.is_empty() { - return (None, AllNamesFlags::empty()); - } - match &args[0].node { - ExprKind::List { elts, .. } - | ExprKind::Set { elts, .. } - | ExprKind::Tuple { elts, .. } => { - return (Some(elts), AllNamesFlags::empty()); - } - ExprKind::ListComp { .. } - | ExprKind::SetComp { .. } - | ExprKind::GeneratorExp { .. } => { - // Allow comprehensions, even though we can't statically analyze - // them. - return (None, AllNamesFlags::empty()); - } - _ => {} - } - } - } - } - _ => {} - } - (None, AllNamesFlags::INVALID_FORMAT) - } - - let mut names: Vec = vec![]; - let mut flags = AllNamesFlags::empty(); - - // Grab the existing bound __all__ values. - if let StmtKind::AugAssign { .. } = &stmt.node { - if let Some(index) = scope.get("__all__") { - if let BindingKind::Export(Export { names: existing }) = &ctx.bindings[*index].kind { - names.extend_from_slice(existing); - } - } - } - - if let Some(value) = match &stmt.node { - StmtKind::Assign { value, .. } => Some(value), - StmtKind::AnnAssign { value, .. } => value.as_ref(), - StmtKind::AugAssign { value, .. } => Some(value), - _ => None, - } { - if let ExprKind::BinOp { left, right, .. } = &value.node { - 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(ctx, current_right); - flags |= new_flags; - if let Some(elts) = elts { - add_to_names(&mut names, elts, &mut flags); - } - - // Process the left side, which can be a "real" value or the "rest" of the - // binary operation. - if let ExprKind::BinOp { left, right, .. } = ¤t_left.node { - current_left = left; - current_right = right; - } else { - let (elts, new_flags) = extract_elts(ctx, current_left); - flags |= new_flags; - if let Some(elts) = elts { - add_to_names(&mut names, elts, &mut flags); - } - break; - } - } - } else { - let (elts, new_flags) = extract_elts(ctx, value); - flags |= new_flags; - if let Some(elts) = elts { - add_to_names(&mut names, elts, &mut flags); - } - } - } - - (names, flags) -} - #[derive(Default)] struct GlobalVisitor<'a> { globals: FxHashMap<&'a str, &'a Stmt>,