diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs index 6f02558e14..778d8fe81f 100644 --- a/crates/ruff/src/checkers/ast/mod.rs +++ b/crates/ruff/src/checkers/ast/mod.rs @@ -1696,11 +1696,7 @@ impl<'a> Checker<'a> { return; } - if self - .semantic - .expr_ancestors() - .any(|expr| expr.is_named_expr_expr()) - { + if self.semantic.expr_ancestors().any(Expr::is_named_expr_expr) { self.add_binding( id, expr.range(), diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs index dd60425922..d12e1e6d28 100644 --- a/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs +++ b/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs @@ -248,7 +248,6 @@ pub(crate) fn unittest_assertion( // the assertion is part of a larger expression. if checker.semantic().stmt().is_expr_stmt() && checker.semantic().expr_parent().is_none() - && !checker.semantic().scope().kind.is_lambda() && !checker.indexer().comment_ranges().intersects(expr.range()) { if let Ok(stmt) = unittest_assert.generate_assert(args, keywords) { diff --git a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs index 67f27c51f7..2282ede228 100644 --- a/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs +++ b/crates/ruff/src/rules/pandas_vet/rules/inplace_argument.rs @@ -84,14 +84,9 @@ pub(crate) fn inplace_argument(checker: &mut Checker, call: &ast::ExprCall) { // the star argument _doesn't_ contain an override). // 2. The call is part of a larger expression (we're converting an expression to a // statement, and expressions can't contain statements). - // 3. The call is in a lambda (we can't assign to a variable in a lambda). This - // should be unnecessary, as lambdas are expressions, and so (2) should apply, - // but we don't currently restore expression stacks when parsing deferred nodes, - // and so the parent is lost. if !seen_star && checker.semantic().stmt().is_expr_stmt() && checker.semantic().expr_parent().is_none() - && !checker.semantic().scope().kind.is_lambda() { if let Some(fix) = convert_inplace_argument_to_assignment( call, diff --git a/crates/ruff/src/rules/pylint/rules/compare_to_empty_string.rs b/crates/ruff/src/rules/pylint/rules/compare_to_empty_string.rs index f784489f60..bb2c4d9049 100644 --- a/crates/ruff/src/rules/pylint/rules/compare_to_empty_string.rs +++ b/crates/ruff/src/rules/pylint/rules/compare_to_empty_string.rs @@ -68,7 +68,7 @@ pub(crate) fn compare_to_empty_string( if checker .semantic() .expr_ancestors() - .any(|parent| parent.is_subscript_expr()) + .any(Expr::is_subscript_expr) { return; } diff --git a/crates/ruff_python_semantic/src/model.rs b/crates/ruff_python_semantic/src/model.rs index bede85c5fd..e66e1c211e 100644 --- a/crates/ruff_python_semantic/src/model.rs +++ b/crates/ruff_python_semantic/src/model.rs @@ -36,8 +36,11 @@ pub struct SemanticModel<'a> { /// The identifier of the current statement. stmt_id: Option, - /// Stack of current expressions. - exprs: Vec<&'a Expr>, + /// Stack of all visited expressions. + exprs: Nodes<'a, Expr>, + + /// The identifier of the current expression. + expr_id: Option, /// Stack of all scopes, along with the identifier of the current scope. pub scopes: Scopes<'a>, @@ -131,7 +134,8 @@ impl<'a> SemanticModel<'a> { module_path: module.path(), stmts: Nodes::::default(), stmt_id: None, - exprs: Vec::default(), + exprs: Nodes::::default(), + expr_id: None, scopes: Scopes::default(), scope_id: ScopeId::global(), definitions: Definitions::for_module(module), @@ -769,16 +773,15 @@ impl<'a> SemanticModel<'a> { self.stmt_id = self.stmts.parent_id(node_id); } - /// Push an [`Expr`] onto the stack. + /// Push a [`Expr`] onto the stack. pub fn push_expr(&mut self, expr: &'a Expr) { - self.exprs.push(expr); + self.expr_id = Some(self.exprs.insert(expr, self.expr_id)); } /// Pop the current [`Expr`] off the stack. pub fn pop_expr(&mut self) { - self.exprs - .pop() - .expect("Attempted to pop without expression"); + let node_id = self.expr_id.expect("Attempted to pop without expression"); + self.expr_id = self.exprs.parent_id(node_id); } /// Push a [`Scope`] with the given [`ScopeKind`] onto the stack. @@ -822,22 +825,32 @@ impl<'a> SemanticModel<'a> { /// Return the current `Expr`. pub fn expr(&self) -> Option<&'a Expr> { - self.exprs.iter().last().copied() + let node_id = self.expr_id?; + Some(self.exprs[node_id]) } - /// Return the parent `Expr` of the current `Expr`. + /// Return the parent `Expr` of the current `Expr`, if any. pub fn expr_parent(&self) -> Option<&'a Expr> { - self.exprs.iter().rev().nth(1).copied() + self.expr_ancestors().next() } - /// Return the grandparent `Expr` of the current `Expr`. + /// Return the grandparent `Expr` of the current `Expr`, if any. pub fn expr_grandparent(&self) -> Option<&'a Expr> { - self.exprs.iter().rev().nth(2).copied() + self.expr_ancestors().nth(1) } /// Return an [`Iterator`] over the current `Expr` parents. - pub fn expr_ancestors(&self) -> impl Iterator { - self.exprs.iter().rev().skip(1) + pub fn expr_ancestors(&self) -> impl Iterator + '_ { + self.expr_id + .iter() + .copied() + .flat_map(|id| { + self.exprs + .ancestor_ids(id) + .skip(1) + .map(|id| &self.exprs[id]) + }) + .copied() } /// Returns a reference to the global scope @@ -1036,6 +1049,7 @@ impl<'a> SemanticModel<'a> { Snapshot { scope_id: self.scope_id, stmt_id: self.stmt_id, + expr_id: self.expr_id, definition_id: self.definition_id, flags: self.flags, } @@ -1046,11 +1060,13 @@ impl<'a> SemanticModel<'a> { let Snapshot { scope_id, stmt_id, + expr_id, definition_id, flags, } = snapshot; self.scope_id = scope_id; self.stmt_id = stmt_id; + self.expr_id = expr_id; self.definition_id = definition_id; self.flags = flags; } @@ -1464,6 +1480,7 @@ impl SemanticModelFlags { pub struct Snapshot { scope_id: ScopeId, stmt_id: Option, + expr_id: Option, definition_id: DefinitionId, flags: SemanticModelFlags, }