Python
/
ruff
mirror of https://github.com/astral-sh/ruff
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Remove suite body tracking from `SemanticModel` (#5848) 

## Summary

The `SemanticModel` currently stores the "body" of a given `Suite`,
along with the current statement index. This is used to support "next
sibling" queries, but we only use this in exactly one place -- the rule
that simplifies constructs like this to `any` or `all`:

```python
for x in y:
    if x == 0:
        return True
return False
```

Instead of tracking the state, we can just do a (slightly more
expensive) traversal, by finding the node within its parent and
returning the next node in the body.

Note that we'll only have to do this extremely rarely -- namely, for
functions that contain something like:

```python
for x in y:
    if x == 0:
        return True
```
This commit is contained in:
Charlie Marsh 2023-07-18 18:58:31 -04:00 committed by GitHub
parent a93254f026
commit 2d505e2b04
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 340 additions and 250 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
crates/ruff/src/checkers/ast/mod.rs
View File

@ -1404,11 +1404,7 @@ where
} }
if stmt.is_for_stmt() { if stmt.is_for_stmt() {
if self.enabled(Rule::ReimplementedBuiltin) { if self.enabled(Rule::ReimplementedBuiltin) {
flake8_simplify::rules::convert_for_loop_to_any_all( flake8_simplify::rules::convert_for_loop_to_any_all(self, stmt);
self,
stmt,
self.semantic.sibling_stmt(),
);
} }
if self.enabled(Rule::InDictKeys) { if self.enabled(Rule::InDictKeys) {
flake8_simplify::rules::key_in_dict_for(self, target, iter); flake8_simplify::rules::key_in_dict_for(self, target, iter);
@ -4237,21 +4233,10 @@ where
flake8_pie::rules::no_unnecessary_pass(self, body); flake8_pie::rules::no_unnecessary_pass(self, body);
} }
// Step 2: Binding
let prev_body = self.semantic.body;
let prev_body_index = self.semantic.body_index;
self.semantic.body = body;
self.semantic.body_index = 0;
// Step 3: Traversal // Step 3: Traversal
for stmt in body { for stmt in body {
self.visit_stmt(stmt); self.visit_stmt(stmt);
self.semantic.body_index += 1;
} }
// Step 4: Clean-up
self.semantic.body = prev_body;
self.semantic.body_index = prev_body_index;
} }
} }

448
crates/ruff/src/rules/flake8_simplify/rules/reimplemented_builtin.rs
View File

@ -1,4 +1,4 @@
use ruff_text_size::{TextRange, TextSize}; use ruff_text_size::TextRange;
use rustpython_parser::ast::{ use rustpython_parser::ast::{
self, CmpOp, Comprehension, Constant, Expr, ExprContext, Ranged, Stmt, UnaryOp, self, CmpOp, Comprehension, Constant, Expr, ExprContext, Ranged, Stmt, UnaryOp,
}; };
@ -7,10 +7,11 @@ use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
use ruff_macros::{derive_message_formats, violation}; use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::any_over_expr; use ruff_python_ast::helpers::any_over_expr;
use ruff_python_ast::source_code::Generator; use ruff_python_ast::source_code::Generator;
use ruff_python_ast::traversal;
use crate::checkers::ast::Checker; use crate::checkers::ast::Checker;
use crate::line_width::LineWidth; use crate::line_width::LineWidth;
use crate::registry::{AsRule, Rule}; use crate::registry::AsRule;
/// ## What it does /// ## What it does
/// Checks for `for` loops that can be replaced with a builtin function, like /// Checks for `for` loops that can be replaced with a builtin function, like
@ -38,7 +39,7 @@ use crate::registry::{AsRule, Rule};
/// - [Python documentation: `all`](https://docs.python.org/3/library/functions.html#all) /// - [Python documentation: `all`](https://docs.python.org/3/library/functions.html#all)
#[violation] #[violation]
pub struct ReimplementedBuiltin { pub struct ReimplementedBuiltin {
repl: String, replacement: String,
} }
impl Violation for ReimplementedBuiltin { impl Violation for ReimplementedBuiltin {
@ -46,200 +47,222 @@ impl Violation for ReimplementedBuiltin {
#[derive_message_formats] #[derive_message_formats]
fn message(&self) -> String { fn message(&self) -> String {
let ReimplementedBuiltin { repl } = self; let ReimplementedBuiltin { replacement } = self;
format!("Use `{repl}` instead of `for` loop") format!("Use `{replacement}` instead of `for` loop")
} }
fn autofix_title(&self) -> Option<String> { fn autofix_title(&self) -> Option<String> {
let ReimplementedBuiltin { repl } = self; let ReimplementedBuiltin { replacement } = self;
Some(format!("Replace with `{repl}`")) Some(format!("Replace with `{replacement}`"))
} }
} }
/// SIM110, SIM111 /// SIM110, SIM111
pub(crate) fn convert_for_loop_to_any_all( pub(crate) fn convert_for_loop_to_any_all(checker: &mut Checker, stmt: &Stmt) {
checker: &mut Checker, if !checker.semantic().scope().kind.is_any_function() {
stmt: &Stmt, return;
sibling: Option<&Stmt>, }
) {
// There are two cases to consider: // The `for` loop itself must consist of an `if` with a `return`.
let Some(loop_) = match_loop(stmt) else {
return;
};
// Afterwards, there are two cases to consider:
// - `for` loop with an `else: return True` or `else: return False`. // - `for` loop with an `else: return True` or `else: return False`.
// - `for` loop followed by `return True` or `return False` // - `for` loop followed by `return True` or `return False`.
if let Some(loop_info) = return_values_for_else(stmt) let Some(terminal) = match_else_return(stmt).or_else(|| {
.or_else(|| sibling.and_then(|sibling| return_values_for_siblings(stmt, sibling))) let parent = checker.semantic().stmt_parent()?;
{ let suite = traversal::suite(stmt, parent)?;
// Check if loop_info.target, loop_info.iter, or loop_info.test contains `await`. let sibling = traversal::next_sibling(stmt, suite)?;
if contains_await(loop_info.target) match_sibling_return(stmt, sibling)
|| contains_await(loop_info.iter) }) else {
|| contains_await(loop_info.test) return;
{ };
return;
}
if loop_info.return_value && !loop_info.next_return_value {
if checker.enabled(Rule::ReimplementedBuiltin) {
let contents = return_stmt(
"any",
loop_info.test,
loop_info.target,
loop_info.iter,
checker.generator(),
);
// Don't flag if the resulting expression would exceed the maximum line length. // Check if any of the expressions contain an `await` expression.
let line_start = checker.locator.line_start(stmt.start()); if contains_await(loop_.target) || contains_await(loop_.iter) || contains_await(loop_.test) {
if LineWidth::new(checker.settings.tab_size) return;
.add_str(&checker.locator.contents()[TextRange::new(line_start, stmt.start())]) }
.add_str(&contents)
> checker.settings.line_length
{
return;
}
let mut diagnostic = Diagnostic::new( match (loop_.return_value, terminal.return_value) {
ReimplementedBuiltin { // Replace with `any`.
repl: contents.clone(), (true, false) => {
}, let contents = return_stmt(
TextRange::new(stmt.start(), loop_info.terminal), "any",
); loop_.test,
if checker.patch(diagnostic.kind.rule()) && checker.semantic().is_builtin("any") { loop_.target,
diagnostic.set_fix(Fix::suggested(Edit::replacement( loop_.iter,
contents, checker.generator(),
stmt.start(), );
loop_info.terminal,
))); // Don't flag if the resulting expression would exceed the maximum line length.
} let line_start = checker.locator.line_start(stmt.start());
checker.diagnostics.push(diagnostic); if LineWidth::new(checker.settings.tab_size)
.add_str(&checker.locator.contents()[TextRange::new(line_start, stmt.start())])
.add_str(&contents)
> checker.settings.line_length
{
return;
} }
}
if !loop_info.return_value && loop_info.next_return_value { let mut diagnostic = Diagnostic::new(
if checker.enabled(Rule::ReimplementedBuiltin) { ReimplementedBuiltin {
// Invert the condition. replacement: contents.to_string(),
let test = { },
if let Expr::UnaryOp(ast::ExprUnaryOp { TextRange::new(stmt.start(), terminal.stmt.end()),
op: UnaryOp::Not, );
operand, if checker.patch(diagnostic.kind.rule()) && checker.semantic().is_builtin("any") {
range: _, diagnostic.set_fix(Fix::suggested(Edit::replacement(
}) = &loop_info.test contents,
{ stmt.start(),
*operand.clone() terminal.stmt.end(),
} else if let Expr::Compare(ast::ExprCompare { )));
left, }
ops, checker.diagnostics.push(diagnostic);
comparators, }
range: _, // Replace with `all`.
}) = &loop_info.test (false, true) => {
{ // Invert the condition.
if let ([op], [comparator]) = (ops.as_slice(), comparators.as_slice()) { let test = {
let op = match op { if let Expr::UnaryOp(ast::ExprUnaryOp {
CmpOp::Eq => CmpOp::NotEq, op: UnaryOp::Not,
CmpOp::NotEq => CmpOp::Eq, operand,
CmpOp::Lt => CmpOp::GtE, range: _,
CmpOp::LtE => CmpOp::Gt, }) = &loop_.test
CmpOp::Gt => CmpOp::LtE, {
CmpOp::GtE => CmpOp::Lt, *operand.clone()
CmpOp::Is => CmpOp::IsNot, } else if let Expr::Compare(ast::ExprCompare {
CmpOp::IsNot => CmpOp::Is, left,
CmpOp::In => CmpOp::NotIn, ops,
CmpOp::NotIn => CmpOp::In, comparators,
}; range: _,
let node = ast::ExprCompare { }) = &loop_.test
left: left.clone(), {
ops: vec![op], if let ([op], [comparator]) = (ops.as_slice(), comparators.as_slice()) {
comparators: vec![comparator.clone()], let op = match op {
range: TextRange::default(), CmpOp::Eq => CmpOp::NotEq,
}; CmpOp::NotEq => CmpOp::Eq,
node.into() CmpOp::Lt => CmpOp::GtE,
} else { CmpOp::LtE => CmpOp::Gt,
let node = ast::ExprUnaryOp { CmpOp::Gt => CmpOp::LtE,
op: UnaryOp::Not, CmpOp::GtE => CmpOp::Lt,
operand: Box::new(loop_info.test.clone()), CmpOp::Is => CmpOp::IsNot,
range: TextRange::default(), CmpOp::IsNot => CmpOp::Is,
}; CmpOp::In => CmpOp::NotIn,
node.into() CmpOp::NotIn => CmpOp::In,
} };
let node = ast::ExprCompare {
left: left.clone(),
ops: vec![op],
comparators: vec![comparator.clone()],
range: TextRange::default(),
};
node.into()
} else { } else {
let node = ast::ExprUnaryOp { let node = ast::ExprUnaryOp {
op: UnaryOp::Not, op: UnaryOp::Not,
operand: Box::new(loop_info.test.clone()), operand: Box::new(loop_.test.clone()),
range: TextRange::default(), range: TextRange::default(),
}; };
node.into() node.into()
} }
}; } else {
let contents = return_stmt( let node = ast::ExprUnaryOp {
"all", op: UnaryOp::Not,
&test, operand: Box::new(loop_.test.clone()),
loop_info.target, range: TextRange::default(),
loop_info.iter, };
checker.generator(), node.into()
);
// Don't flag if the resulting expression would exceed the maximum line length.
let line_start = checker.locator.line_start(stmt.start());
if LineWidth::new(checker.settings.tab_size)
.add_str(&checker.locator.contents()[TextRange::new(line_start, stmt.start())])
.add_str(&contents)
> checker.settings.line_length
{
return;
} }
};
let contents = return_stmt("all", &test, loop_.target, loop_.iter, checker.generator());
let mut diagnostic = Diagnostic::new( // Don't flag if the resulting expression would exceed the maximum line length.
ReimplementedBuiltin { let line_start = checker.locator.line_start(stmt.start());
repl: contents.clone(), if LineWidth::new(checker.settings.tab_size)
}, .add_str(&checker.locator.contents()[TextRange::new(line_start, stmt.start())])
TextRange::new(stmt.start(), loop_info.terminal), .add_str(&contents)
); > checker.settings.line_length
if checker.patch(diagnostic.kind.rule()) && checker.semantic().is_builtin("all") { {
diagnostic.set_fix(Fix::suggested(Edit::replacement( return;
contents,
stmt.start(),
loop_info.terminal,
)));
}
checker.diagnostics.push(diagnostic);
} }
let mut diagnostic = Diagnostic::new(
ReimplementedBuiltin {
replacement: contents.to_string(),
},
TextRange::new(stmt.start(), terminal.stmt.end()),
);
if checker.patch(diagnostic.kind.rule()) && checker.semantic().is_builtin("all") {
diagnostic.set_fix(Fix::suggested(Edit::replacement(
contents,
stmt.start(),
terminal.stmt.end(),
)));
}
checker.diagnostics.push(diagnostic);
} }
_ => {}
} }
} }
/// Represents a `for` loop with a conditional `return`, like:
/// ```python
/// for x in y:
/// if x == 0:
/// return True
/// ```
#[derive(Debug)]
struct Loop<'a> { struct Loop<'a> {
/// The `return` value of the loop.
return_value: bool, return_value: bool,
next_return_value: bool, /// The test condition in the loop.
test: &'a Expr, test: &'a Expr,
/// The target of the loop.
target: &'a Expr, target: &'a Expr,
/// The iterator of the loop.
iter: &'a Expr, iter: &'a Expr,
terminal: TextSize,
} }
/// Extract the returned boolean values a `Stmt::For` with an `else` body. /// Represents a `return` statement following a `for` loop, like:
fn return_values_for_else(stmt: &Stmt) -> Option<Loop> { /// ```python
/// for x in y:
/// if x == 0:
/// return True
/// return False
/// ```
///
/// Or:
/// ```python
/// for x in y:
/// if x == 0:
/// return True
/// else:
/// return False
/// ```
#[derive(Debug)]
struct Terminal<'a> {
return_value: bool,
stmt: &'a Stmt,
}
fn match_loop(stmt: &Stmt) -> Option<Loop> {
let Stmt::For(ast::StmtFor { let Stmt::For(ast::StmtFor {
body, body, target, iter, ..
target,
iter,
orelse,
..
}) = stmt }) = stmt
else { else {
return None; return None;
}; };
// The loop itself should contain a single `if` statement, with an `else` // The loop itself should contain a single `if` statement, with a single `return` statement in
// containing a single `return True` or `return False`. // the body.
if body.len() != 1 { let [Stmt::If(ast::StmtIf {
return None;
}
if orelse.len() != 1 {
return None;
}
let Stmt::If(ast::StmtIf {
body: nested_body, body: nested_body,
test: nested_test, test: nested_test,
elif_else_clauses: nested_elif_else_clauses, elif_else_clauses: nested_elif_else_clauses,
range: _, range: _,
}) = &body[0] })] = body.as_slice()
else { else {
return None; return None;
}; };
@ -263,15 +286,35 @@ fn return_values_for_else(stmt: &Stmt) -> Option<Loop> {
return None; return None;
}; };
// The `else` block has to contain a single `return True` or `return False`. Some(Loop {
let Stmt::Return(ast::StmtReturn { return_value: *value,
value: next_value, test: nested_test,
range: _, target,
}) = &orelse[0] iter,
else { })
}
/// If a `Stmt::For` contains an `else` with a single boolean `return`, return the [`Terminal`]
/// representing that `return`.
///
/// For example, matches the `return` in:
/// ```python
/// for x in y:
/// if x == 0:
/// return True
/// return False
/// ```
fn match_else_return(stmt: &Stmt) -> Option<Terminal> {
let Stmt::For(ast::StmtFor { orelse, .. }) = stmt else {
return None; return None;
}; };
let Some(next_value) = next_value else {
// The `else` block has to contain a single `return True` or `return False`.
let [Stmt::Return(ast::StmtReturn {
value: Some(next_value),
range: _,
})] = orelse.as_slice()
else {
return None; return None;
}; };
let Expr::Constant(ast::ExprConstant { let Expr::Constant(ast::ExprConstant {
@ -282,78 +325,41 @@ fn return_values_for_else(stmt: &Stmt) -> Option<Loop> {
return None; return None;
}; };
Some(Loop { Some(Terminal {
return_value: *value, return_value: *next_value,
next_return_value: *next_value, stmt,
test: nested_test,
target,
iter,
terminal: stmt.end(),
}) })
} }
/// Extract the returned boolean values from subsequent `Stmt::For` and /// If a `Stmt::For` is followed by a boolean `return`, return the [`Terminal`] representing that
/// `Stmt::Return` statements, or `None`. /// `return`.
fn return_values_for_siblings<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option<Loop<'a>> { ///
let Stmt::For(ast::StmtFor { /// For example, matches the `return` in:
body, /// ```python
target, /// for x in y:
iter, /// if x == 0:
orelse, /// return True
.. /// else:
}) = stmt /// return False
else { /// ```
fn match_sibling_return<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option<Terminal<'a>> {
let Stmt::For(ast::StmtFor { orelse, .. }) = stmt else {
return None; return None;
}; };
// The loop itself should contain a single `if` statement, with a single `return // The loop itself shouldn't have an `else` block.
// True` or `return False`.
if body.len() != 1 {
return None;
}
if !orelse.is_empty() { if !orelse.is_empty() {
return None; return None;
} }
let Stmt::If(ast::StmtIf {
body: nested_body,
test: nested_test,
elif_else_clauses: nested_elif_else_clauses,
range: _,
}) = &body[0]
else {
return None;
};
if nested_body.len() != 1 {
return None;
}
if !nested_elif_else_clauses.is_empty() {
return None;
}
let Stmt::Return(ast::StmtReturn { value, range: _ }) = &nested_body[0] else {
return None;
};
let Some(value) = value else {
return None;
};
let Expr::Constant(ast::ExprConstant {
value: Constant::Bool(value),
..
}) = value.as_ref()
else {
return None;
};
// The next statement has to be a `return True` or `return False`. // The next statement has to be a `return True` or `return False`.
let Stmt::Return(ast::StmtReturn { let Stmt::Return(ast::StmtReturn {
value: next_value, value: Some(next_value),
range: _, range: _,
}) = &sibling }) = &sibling
else { else {
return None; return None;
}; };
let Some(next_value) = next_value else {
return None;
};
let Expr::Constant(ast::ExprConstant { let Expr::Constant(ast::ExprConstant {
value: Constant::Bool(next_value), value: Constant::Bool(next_value),
.. ..
@ -362,13 +368,9 @@ fn return_values_for_siblings<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option<L
return None; return None;
}; };
Some(Loop { Some(Terminal {
return_value: *value, return_value: *next_value,
next_return_value: *next_value, stmt: sibling,
test: nested_test,
target,
iter,
terminal: sibling.end(),
}) })
} }

1
crates/ruff_python_ast/src/lib.rs
View File

@ -15,6 +15,7 @@ pub mod statement_visitor;
pub mod stmt_if; pub mod stmt_if;
pub mod str; pub mod str;
pub mod token_kind; pub mod token_kind;
pub mod traversal;
pub mod types; pub mod types;
pub mod typing; pub mod typing;
pub mod visitor; pub mod visitor;

113
crates/ruff_python_ast/src/traversal.rs Normal file
View File

@ -0,0 +1,113 @@
//! Utilities for manually traversing a Python AST.
use rustpython_ast::{ExceptHandler, Stmt, Suite};
use rustpython_parser::ast;
/// Given a [`Stmt`] and its parent, return the [`Suite`] that contains the [`Stmt`].
pub fn suite<'a>(stmt: &'a Stmt, parent: &'a Stmt) -> Option<&'a Suite> {
match parent {
Stmt::FunctionDef(ast::StmtFunctionDef { body, .. }) => Some(body),
Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef { body, .. }) => Some(body),
Stmt::ClassDef(ast::StmtClassDef { body, .. }) => Some(body),
Stmt::For(ast::StmtFor { body, orelse, .. }) => {
if body.contains(stmt) {
Some(body)
} else if orelse.contains(stmt) {
Some(orelse)
} else {
None
}
}
Stmt::AsyncFor(ast::StmtAsyncFor { body, orelse, .. }) => {
if body.contains(stmt) {
Some(body)
} else if orelse.contains(stmt) {
Some(orelse)
} else {
None
}
}
Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
if body.contains(stmt) {
Some(body)
} else if orelse.contains(stmt) {
Some(orelse)
} else {
None
}
}
Stmt::If(ast::StmtIf {
body,
elif_else_clauses,
..
}) => {
if body.contains(stmt) {
Some(body)
} else {
elif_else_clauses
.iter()
.map(|elif_else_clause| &elif_else_clause.body)
.find(|body| body.contains(stmt))
}
}
Stmt::With(ast::StmtWith { body, .. }) => Some(body),
Stmt::AsyncWith(ast::StmtAsyncWith { body, .. }) => Some(body),
Stmt::Match(ast::StmtMatch { cases, .. }) => cases
.iter()
.map(|case| &case.body)
.find(|body| body.contains(stmt)),
Stmt::Try(ast::StmtTry {
body,
handlers,
orelse,
finalbody,
..
}) => {
if body.contains(stmt) {
Some(body)
} else if orelse.contains(stmt) {
Some(orelse)
} else if finalbody.contains(stmt) {
Some(finalbody)
} else {
handlers
.iter()
.filter_map(ExceptHandler::as_except_handler)
.map(|handler| &handler.body)
.find(|body| body.contains(stmt))
}
}
Stmt::TryStar(ast::StmtTryStar {
body,
handlers,
orelse,
finalbody,
..
}) => {
if body.contains(stmt) {
Some(body)
} else if orelse.contains(stmt) {
Some(orelse)
} else if finalbody.contains(stmt) {
Some(finalbody)
} else {
handlers
.iter()
.filter_map(ExceptHandler::as_except_handler)
.map(|handler| &handler.body)
.find(|body| body.contains(stmt))
}
}
_ => None,
}
}
/// Given a [`Stmt`] and its containing [`Suite`], return the next [`Stmt`] in the [`Suite`].
pub fn next_sibling<'a>(stmt: &'a Stmt, suite: &'a Suite) -> Option<&'a Stmt> {
let mut iter = suite.iter();
while let Some(sibling) = iter.next() {
if sibling == stmt {
return iter.next();
}
}
None
}

11
crates/ruff_python_semantic/src/model.rs
View File

@ -108,10 +108,6 @@ pub struct SemanticModel<'a> {
/// by way of the `global x` statement. /// by way of the `global x` statement.
rebinding_scopes: HashMap<BindingId, Vec<ScopeId>, BuildNoHashHasher<BindingId>>, rebinding_scopes: HashMap<BindingId, Vec<ScopeId>, BuildNoHashHasher<BindingId>>,
/// Body iteration; used to peek at siblings.
pub body: &'a [Stmt],
pub body_index: usize,
/// Flags for the semantic model. /// Flags for the semantic model.
pub flags: SemanticModelFlags, pub flags: SemanticModelFlags,
@ -137,8 +133,6 @@ impl<'a> SemanticModel<'a> {
shadowed_bindings: IntMap::default(), shadowed_bindings: IntMap::default(),
delayed_annotations: IntMap::default(), delayed_annotations: IntMap::default(),
rebinding_scopes: IntMap::default(), rebinding_scopes: IntMap::default(),
body: &[],
body_index: 0,
flags: SemanticModelFlags::new(path), flags: SemanticModelFlags::new(path),
handled_exceptions: Vec::default(), handled_exceptions: Vec::default(),
} }
@ -757,11 +751,6 @@ impl<'a> SemanticModel<'a> {
self.exprs.iter().rev().skip(1) self.exprs.iter().rev().skip(1)
} }
/// Return the `Stmt` that immediately follows the current `Stmt`, if any.
pub fn sibling_stmt(&self) -> Option<&'a Stmt> {
self.body.get(self.body_index + 1)
}
/// Returns a reference to the global scope /// Returns a reference to the global scope
pub fn global_scope(&self) -> &Scope<'a> { pub fn global_scope(&self) -> &Scope<'a> {
self.scopes.global() self.scopes.global()