mirror of
https://github.com/astral-sh/ruff.git
synced 2025-07-25 05:53:51 +00:00
Combine operations.rs and helpers.rs (#3841)
This commit is contained in:
Charlie Marsh
GitHub
parent
2fbc620ad3
commit
6d80c79bac
5 changed files with 281 additions and 298 deletions
|
|
@ -6,8 +6,8 @@ use once_cell::sync::Lazy;
|
|||
use regex::Regex;
|
||||
use rustc_hash::{FxHashMap, FxHashSet};
|
||||
use rustpython_parser::ast::{
|
||||
Arguments, Constant, Excepthandler, ExcepthandlerKind, Expr, ExprKind, Keyword, KeywordData,
|
||||
Located, Location, MatchCase, Pattern, PatternKind, Stmt, StmtKind,
|
||||
Arguments, Cmpop, Constant, Excepthandler, ExcepthandlerKind, Expr, ExprKind, Keyword,
|
||||
KeywordData, Located, Location, MatchCase, Pattern, PatternKind, Stmt, StmtKind,
|
||||
};
|
||||
use rustpython_parser::{lexer, Mode, Tok};
|
||||
use smallvec::{smallvec, SmallVec};
|
||||
|
|
@ -852,6 +852,38 @@ where
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct GlobalStatementVisitor<'a> {
|
||||
globals: FxHashMap<&'a str, &'a Stmt>,
|
||||
}
|
||||
|
||||
impl<'a> Visitor<'a> for GlobalStatementVisitor<'a> {
|
||||
fn visit_stmt(&mut self, stmt: &'a Stmt) {
|
||||
match &stmt.node {
|
||||
StmtKind::Global { names } => {
|
||||
for name in names {
|
||||
self.globals.insert(name, stmt);
|
||||
}
|
||||
}
|
||||
StmtKind::FunctionDef { .. }
|
||||
| StmtKind::AsyncFunctionDef { .. }
|
||||
| StmtKind::ClassDef { .. } => {
|
||||
// Don't recurse.
|
||||
}
|
||||
_ => visitor::walk_stmt(self, stmt),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract a map from global name to its last-defining [`Stmt`].
|
||||
pub fn extract_globals(body: &[Stmt]) -> FxHashMap<&str, &Stmt> {
|
||||
let mut visitor = GlobalStatementVisitor::default();
|
||||
for stmt in body {
|
||||
visitor.visit_stmt(stmt);
|
||||
}
|
||||
visitor.globals
|
||||
}
|
||||
|
||||
/// Convert a location within a file (relative to `base`) to an absolute
|
||||
/// position.
|
||||
pub fn to_absolute(relative: Location, base: Location) -> Location {
|
||||
|
|
@ -1231,14 +1263,182 @@ pub fn is_logger_candidate(context: &Context, func: &Expr) -> bool {
|
|||
false
|
||||
}
|
||||
|
||||
/// Check if a node is parent of a conditional branch.
|
||||
pub fn on_conditional_branch<'a>(parents: &mut impl Iterator<Item = &'a Stmt>) -> bool {
|
||||
parents.any(|parent| {
|
||||
if matches!(
|
||||
parent.node,
|
||||
StmtKind::If { .. } | StmtKind::While { .. } | StmtKind::Match { .. }
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
if let StmtKind::Expr { value } = &parent.node {
|
||||
if matches!(value.node, ExprKind::IfExp { .. }) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
})
|
||||
}
|
||||
|
||||
/// Check if a node is in a nested block.
|
||||
pub fn in_nested_block<'a>(mut parents: impl Iterator<Item = &'a Stmt>) -> bool {
|
||||
parents.any(|parent| {
|
||||
matches!(
|
||||
parent.node,
|
||||
StmtKind::Try { .. }
|
||||
| StmtKind::TryStar { .. }
|
||||
| StmtKind::If { .. }
|
||||
| StmtKind::With { .. }
|
||||
| StmtKind::Match { .. }
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Check if a node represents an unpacking assignment.
|
||||
pub fn is_unpacking_assignment(parent: &Stmt, child: &Expr) -> bool {
|
||||
match &parent.node {
|
||||
StmtKind::With { items, .. } => items.iter().any(|item| {
|
||||
if let Some(optional_vars) = &item.optional_vars {
|
||||
if matches!(optional_vars.node, ExprKind::Tuple { .. }) {
|
||||
if any_over_expr(optional_vars, &|expr| expr == child) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}),
|
||||
StmtKind::Assign { targets, value, .. } => {
|
||||
// In `(a, b) = (1, 2)`, `(1, 2)` is the target, and it is a tuple.
|
||||
let value_is_tuple = matches!(
|
||||
&value.node,
|
||||
ExprKind::Set { .. } | ExprKind::List { .. } | ExprKind::Tuple { .. }
|
||||
);
|
||||
// In `(a, b) = coords = (1, 2)`, `(a, b)` and `coords` are the targets, and
|
||||
// `(a, b)` is a tuple. (We use "tuple" as a placeholder for any
|
||||
// unpackable expression.)
|
||||
let targets_are_tuples = targets.iter().all(|item| {
|
||||
matches!(
|
||||
item.node,
|
||||
ExprKind::Set { .. } | ExprKind::List { .. } | ExprKind::Tuple { .. }
|
||||
)
|
||||
});
|
||||
// If we're looking at `a` in `(a, b) = coords = (1, 2)`, then we should
|
||||
// identify that the current expression is in a tuple.
|
||||
let child_in_tuple = targets_are_tuples
|
||||
|| targets.iter().any(|item| {
|
||||
matches!(
|
||||
item.node,
|
||||
ExprKind::Set { .. } | ExprKind::List { .. } | ExprKind::Tuple { .. }
|
||||
) && any_over_expr(item, &|expr| expr == child)
|
||||
});
|
||||
|
||||
// If our child is a tuple, and value is not, it's always an unpacking
|
||||
// expression. Ex) `x, y = tup`
|
||||
if child_in_tuple && !value_is_tuple {
|
||||
return true;
|
||||
}
|
||||
|
||||
// If our child isn't a tuple, but value is, it's never an unpacking expression.
|
||||
// Ex) `coords = (1, 2)`
|
||||
if !child_in_tuple && value_is_tuple {
|
||||
return false;
|
||||
}
|
||||
|
||||
// If our target and the value are both tuples, then it's an unpacking
|
||||
// expression assuming there's at least one non-tuple child.
|
||||
// Ex) Given `(x, y) = coords = 1, 2`, `(x, y)` is considered an unpacking
|
||||
// expression. Ex) Given `(x, y) = (a, b) = 1, 2`, `(x, y)` isn't
|
||||
// considered an unpacking expression.
|
||||
if child_in_tuple && value_is_tuple {
|
||||
return !targets_are_tuples;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub type LocatedCmpop<U = ()> = Located<Cmpop, U>;
|
||||
|
||||
/// Extract all [`Cmpop`] operators from a source code snippet, with appropriate
|
||||
/// ranges.
|
||||
///
|
||||
/// `RustPython` doesn't include line and column information on [`Cmpop`] nodes.
|
||||
/// `CPython` doesn't either. This method iterates over the token stream and
|
||||
/// re-identifies [`Cmpop`] nodes, annotating them with valid ranges.
|
||||
pub fn locate_cmpops(contents: &str) -> Vec<LocatedCmpop> {
|
||||
let mut tok_iter = lexer::lex(contents, Mode::Module).flatten().peekable();
|
||||
let mut ops: Vec<LocatedCmpop> = vec![];
|
||||
let mut count: usize = 0;
|
||||
loop {
|
||||
let Some((start, tok, end)) = tok_iter.next() else {
|
||||
break;
|
||||
};
|
||||
if matches!(tok, Tok::Lpar) {
|
||||
count += 1;
|
||||
continue;
|
||||
} else if matches!(tok, Tok::Rpar) {
|
||||
count -= 1;
|
||||
continue;
|
||||
}
|
||||
if count == 0 {
|
||||
match tok {
|
||||
Tok::Not => {
|
||||
if let Some((_, _, end)) =
|
||||
tok_iter.next_if(|(_, tok, _)| matches!(tok, Tok::In))
|
||||
{
|
||||
ops.push(LocatedCmpop::new(start, end, Cmpop::NotIn));
|
||||
}
|
||||
}
|
||||
Tok::In => {
|
||||
ops.push(LocatedCmpop::new(start, end, Cmpop::In));
|
||||
}
|
||||
Tok::Is => {
|
||||
let op = if let Some((_, _, end)) =
|
||||
tok_iter.next_if(|(_, tok, _)| matches!(tok, Tok::Not))
|
||||
{
|
||||
LocatedCmpop::new(start, end, Cmpop::IsNot)
|
||||
} else {
|
||||
LocatedCmpop::new(start, end, Cmpop::Is)
|
||||
};
|
||||
ops.push(op);
|
||||
}
|
||||
Tok::NotEqual => {
|
||||
ops.push(LocatedCmpop::new(start, end, Cmpop::NotEq));
|
||||
}
|
||||
Tok::EqEqual => {
|
||||
ops.push(LocatedCmpop::new(start, end, Cmpop::Eq));
|
||||
}
|
||||
Tok::GreaterEqual => {
|
||||
ops.push(LocatedCmpop::new(start, end, Cmpop::GtE));
|
||||
}
|
||||
Tok::Greater => {
|
||||
ops.push(LocatedCmpop::new(start, end, Cmpop::Gt));
|
||||
}
|
||||
Tok::LessEqual => {
|
||||
ops.push(LocatedCmpop::new(start, end, Cmpop::LtE));
|
||||
}
|
||||
Tok::Less => {
|
||||
ops.push(LocatedCmpop::new(start, end, Cmpop::Lt));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
ops
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use anyhow::Result;
|
||||
use rustpython_parser as parser;
|
||||
use rustpython_parser::ast::Location;
|
||||
use rustpython_parser::ast::{Cmpop, Location};
|
||||
|
||||
use crate::helpers::{
|
||||
elif_else_range, else_range, first_colon_range, identifier_range, match_trailing_content,
|
||||
elif_else_range, else_range, first_colon_range, identifier_range, locate_cmpops,
|
||||
match_trailing_content, LocatedCmpop,
|
||||
};
|
||||
use crate::source_code::Locator;
|
||||
use crate::types::Range;
|
||||
|
|
@ -1352,7 +1552,7 @@ class Class():
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn test_else_range() -> Result<()> {
|
||||
fn extract_else_range() -> Result<()> {
|
||||
let contents = r#"
|
||||
for x in y:
|
||||
pass
|
||||
|
|
@ -1372,7 +1572,7 @@ else:
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn test_first_colon_range() {
|
||||
fn extract_first_colon_range() {
|
||||
let contents = "with a: pass";
|
||||
let locator = Locator::new(contents);
|
||||
let range = first_colon_range(
|
||||
|
|
@ -1387,7 +1587,7 @@ else:
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn test_elif_else_range() -> Result<()> {
|
||||
fn extract_elif_else_range() -> Result<()> {
|
||||
let contents = "
|
||||
if a:
|
||||
...
|
||||
|
|
@ -1420,4 +1620,70 @@ else:
|
|||
assert_eq!(range.end_location.column(), 4);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_cmpop_location() {
|
||||
assert_eq!(
|
||||
locate_cmpops("x == 1"),
|
||||
vec![LocatedCmpop::new(
|
||||
Location::new(1, 2),
|
||||
Location::new(1, 4),
|
||||
Cmpop::Eq
|
||||
)]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
locate_cmpops("x != 1"),
|
||||
vec![LocatedCmpop::new(
|
||||
Location::new(1, 2),
|
||||
Location::new(1, 4),
|
||||
Cmpop::NotEq
|
||||
)]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
locate_cmpops("x is 1"),
|
||||
vec![LocatedCmpop::new(
|
||||
Location::new(1, 2),
|
||||
Location::new(1, 4),
|
||||
Cmpop::Is
|
||||
)]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
locate_cmpops("x is not 1"),
|
||||
vec![LocatedCmpop::new(
|
||||
Location::new(1, 2),
|
||||
Location::new(1, 8),
|
||||
Cmpop::IsNot
|
||||
)]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
locate_cmpops("x in 1"),
|
||||
vec![LocatedCmpop::new(
|
||||
Location::new(1, 2),
|
||||
Location::new(1, 4),
|
||||
Cmpop::In
|
||||
)]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
locate_cmpops("x not in 1"),
|
||||
vec![LocatedCmpop::new(
|
||||
Location::new(1, 2),
|
||||
Location::new(1, 8),
|
||||
Cmpop::NotIn
|
||||
)]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
locate_cmpops("x != (1 is not 2)"),
|
||||
vec![LocatedCmpop::new(
|
||||
Location::new(1, 2),
|
||||
Location::new(1, 4),
|
||||
Cmpop::NotEq
|
||||
)]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue