ruff/crates/ruff_python_parser/src/parser/helpers.rs

46 lines
1.8 KiB
Rust

use ruff_python_ast::{self as ast, CmpOp, Expr, ExprContext};
use crate::TokenKind;
/// Set the `ctx` for `Expr::Id`, `Expr::Attribute`, `Expr::Subscript`, `Expr::Starred`,
/// `Expr::Tuple` and `Expr::List`. If `expr` is either `Expr::Tuple` or `Expr::List`,
/// recursively sets the `ctx` for their elements.
pub(super) fn set_expr_ctx(expr: &mut Expr, new_ctx: ExprContext) {
match expr {
Expr::Name(ast::ExprName { ctx, .. })
| Expr::Attribute(ast::ExprAttribute { ctx, .. })
| Expr::Subscript(ast::ExprSubscript { ctx, .. }) => *ctx = new_ctx,
Expr::Starred(ast::ExprStarred { value, ctx, .. }) => {
*ctx = new_ctx;
set_expr_ctx(value, new_ctx);
}
Expr::UnaryOp(ast::ExprUnaryOp { operand, .. }) => {
set_expr_ctx(operand, new_ctx);
}
Expr::List(ast::ExprList { elts, ctx, .. })
| Expr::Tuple(ast::ExprTuple { elts, ctx, .. }) => {
*ctx = new_ctx;
elts.iter_mut()
.for_each(|element| set_expr_ctx(element, new_ctx));
}
_ => {}
}
}
/// Converts a [`TokenKind`] array of size 2 to its correspondent [`CmpOp`].
pub(super) const fn token_kind_to_cmp_op(tokens: [TokenKind; 2]) -> Option<CmpOp> {
Some(match tokens {
[TokenKind::Is, TokenKind::Not] => CmpOp::IsNot,
[TokenKind::Is, _] => CmpOp::Is,
[TokenKind::Not, TokenKind::In] => CmpOp::NotIn,
[TokenKind::In, _] => CmpOp::In,
[TokenKind::EqEqual, _] => CmpOp::Eq,
[TokenKind::NotEqual, _] => CmpOp::NotEq,
[TokenKind::Less, _] => CmpOp::Lt,
[TokenKind::LessEqual, _] => CmpOp::LtE,
[TokenKind::Greater, _] => CmpOp::Gt,
[TokenKind::GreaterEqual, _] => CmpOp::GtE,
_ => return None,
})
}