Skip all bracketed expressions when locating comparison ops (#7740)

Closes https://github.com/astral-sh/ruff/issues/7737.
This commit is contained in:
Charlie Marsh 2023-10-01 10:57:40 -04:00 committed by GitHub
parent 1cf3b5676f
commit 2838f7af98
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 114 additions and 52 deletions

View file

@ -180,66 +180,85 @@ pub fn locate_cmp_ops(expr: &Expr, source: &str) -> Vec<LocatedCmpOp> {
.peekable();
let mut ops: Vec<LocatedCmpOp> = vec![];
let mut count = 0u32;
// Track the bracket depth.
let mut par_count = 0u32;
let mut sqb_count = 0u32;
let mut brace_count = 0u32;
loop {
let Some((tok, range)) = tok_iter.next() else {
break;
};
if matches!(tok, Tok::Lpar) {
count = count.saturating_add(1);
continue;
} else if matches!(tok, Tok::Rpar) {
count = count.saturating_sub(1);
match tok {
Tok::Lpar => {
par_count = par_count.saturating_add(1);
}
Tok::Rpar => {
par_count = par_count.saturating_sub(1);
}
Tok::Lsqb => {
sqb_count = sqb_count.saturating_add(1);
}
Tok::Rsqb => {
sqb_count = sqb_count.saturating_sub(1);
}
Tok::Lbrace => {
brace_count = brace_count.saturating_add(1);
}
Tok::Rbrace => {
brace_count = brace_count.saturating_sub(1);
}
_ => {}
}
if par_count > 0 || sqb_count > 0 || brace_count > 0 {
continue;
}
if count == 0 {
match tok {
Tok::Not => {
if let Some((_, next_range)) =
tok_iter.next_if(|(tok, _)| matches!(tok, Tok::In))
{
ops.push(LocatedCmpOp::new(
TextRange::new(range.start(), next_range.end()),
CmpOp::NotIn,
));
}
match tok {
Tok::Not => {
if let Some((_, next_range)) = tok_iter.next_if(|(tok, _)| tok.is_in()) {
ops.push(LocatedCmpOp::new(
TextRange::new(range.start(), next_range.end()),
CmpOp::NotIn,
));
}
Tok::In => {
ops.push(LocatedCmpOp::new(range, CmpOp::In));
}
Tok::Is => {
let op = if let Some((_, next_range)) =
tok_iter.next_if(|(tok, _)| matches!(tok, Tok::Not))
{
LocatedCmpOp::new(
TextRange::new(range.start(), next_range.end()),
CmpOp::IsNot,
)
} else {
LocatedCmpOp::new(range, CmpOp::Is)
};
ops.push(op);
}
Tok::NotEqual => {
ops.push(LocatedCmpOp::new(range, CmpOp::NotEq));
}
Tok::EqEqual => {
ops.push(LocatedCmpOp::new(range, CmpOp::Eq));
}
Tok::GreaterEqual => {
ops.push(LocatedCmpOp::new(range, CmpOp::GtE));
}
Tok::Greater => {
ops.push(LocatedCmpOp::new(range, CmpOp::Gt));
}
Tok::LessEqual => {
ops.push(LocatedCmpOp::new(range, CmpOp::LtE));
}
Tok::Less => {
ops.push(LocatedCmpOp::new(range, CmpOp::Lt));
}
_ => {}
}
Tok::In => {
ops.push(LocatedCmpOp::new(range, CmpOp::In));
}
Tok::Is => {
let op = if let Some((_, next_range)) = tok_iter.next_if(|(tok, _)| tok.is_not()) {
LocatedCmpOp::new(
TextRange::new(range.start(), next_range.end()),
CmpOp::IsNot,
)
} else {
LocatedCmpOp::new(range, CmpOp::Is)
};
ops.push(op);
}
Tok::NotEqual => {
ops.push(LocatedCmpOp::new(range, CmpOp::NotEq));
}
Tok::EqEqual => {
ops.push(LocatedCmpOp::new(range, CmpOp::Eq));
}
Tok::GreaterEqual => {
ops.push(LocatedCmpOp::new(range, CmpOp::GtE));
}
Tok::Greater => {
ops.push(LocatedCmpOp::new(range, CmpOp::Gt));
}
Tok::LessEqual => {
ops.push(LocatedCmpOp::new(range, CmpOp::LtE));
}
Tok::Less => {
ops.push(LocatedCmpOp::new(range, CmpOp::Lt));
}
_ => {}
}
}
ops