use alloc::fmt;
use core::fmt::Display as _;
use ruff_python_ast as ast;
use ruff_text_size::Ranged;
use rustpython_compiler_core::SourceFile;
use rustpython_literal::escape::{AsciiEscape, UnicodeEscape};
pub(crate) mod precedence {
macro_rules! precedence {
($($op:ident,)*) => {
precedence!(@0, $($op,)*);
};
(@$i:expr, $op1:ident, $($op:ident,)*) => {
pub(crate) const $op1: u8 = $i;
precedence!(@$i + 1, $($op,)*);
};
(@$i:expr,) => {};
}
precedence!(
TUPLE, TEST, OR, AND, NOT, CMP, // "EXPR" =
BOR, BXOR, BAND, SHIFT, ARITH, TERM, FACTOR, POWER, AWAIT, ATOM,
);
pub(crate) const EXPR: u8 = BOR;
}
struct Unparser<'a, 'b, 'c> {
f: &'b mut fmt::Formatter<'a>,
source: &'c SourceFile,
}
impl<'a, 'b, 'c> Unparser<'a, 'b, 'c> {
const fn new(f: &'b mut fmt::Formatter<'a>, source: &'c SourceFile) -> Self {
Self { f, source }
}
fn p(&mut self, s: &str) -> fmt::Result {
self.f.write_str(s)
}
fn p_id(&mut self, s: &ast::Identifier) -> fmt::Result {
self.f.write_str(s.as_str())
}
fn p_if(&mut self, cond: bool, s: &str) -> fmt::Result {
if cond {
self.f.write_str(s)?;
}
Ok(())
}
fn p_delim(&mut self, first: &mut bool, s: &str) -> fmt::Result {
self.p_if(!core::mem::take(first), s)
}
fn write_fmt(&mut self, f: fmt::Arguments<'_>) -> fmt::Result {
self.f.write_fmt(f)
}
fn unparse_expr(&mut self, ast: &ast::Expr, level: u8) -> fmt::Result {
macro_rules! op_prec {
($op_ty:ident, $x:expr, $enu:path, $($var:ident($op:literal, $prec:ident)),*$(,)?) => {
match $x {
$(<$enu>::$var => (op_prec!(@space $op_ty, $op), precedence::$prec),)*
}
};
(@space bin, $op:literal) => {
concat!(" ", $op, " ")
};
(@space un, $op:literal) => {
$op
};
}
macro_rules! group_if {
($lvl:expr, $body:block) => {{
let group = level > $lvl;
self.p_if(group, "(")?;
let ret = $body;
self.p_if(group, ")")?;
ret
}};
}
match &ast {
ast::Expr::BoolOp(ast::ExprBoolOp {
op,
values,
node_index: _,
range: _range,
}) => {
let (op, prec) = op_prec!(bin, op, ast::BoolOp, And("and", AND), Or("or", OR));
group_if!(prec, {
let mut first = true;
for val in values {
self.p_delim(&mut first, op)?;
self.unparse_expr(val, prec + 1)?;
}
})
}
ast::Expr::Named(ast::ExprNamed {
target,
value,
node_index: _,
range: _range,
}) => {
group_if!(precedence::TUPLE, {
self.unparse_expr(target, precedence::ATOM)?;
self.p(" := ")?;
self.unparse_expr(value, precedence::ATOM)?;
})
}
ast::Expr::BinOp(ast::ExprBinOp {
left,
op,
right,
node_index: _,
range: _range,
}) => {
let right_associative = matches!(op, ast::Operator::Pow);
let (op, prec) = op_prec!(
bin,
op,
ast::Operator,
Add("+", ARITH),
Sub("-", ARITH),
Mult("*", TERM),
MatMult("@", TERM),
Div("/", TERM),
Mod("%", TERM),
Pow("**", POWER),
LShift("<<", SHIFT),
RShift(">>", SHIFT),
BitOr("|", BOR),
BitXor("^", BXOR),
BitAnd("&", BAND),
FloorDiv("//", TERM),
);
group_if!(prec, {
self.unparse_expr(left, prec + right_associative as u8)?;
self.p(op)?;
self.unparse_expr(right, prec + !right_associative as u8)?;
})
}
ast::Expr::UnaryOp(ast::ExprUnaryOp {
op,
operand,
node_index: _,
range: _range,
}) => {
let (op, prec) = op_prec!(
un,
op,
ast::UnaryOp,
Invert("~", FACTOR),
Not("not ", NOT),
UAdd("+", FACTOR),
USub("-", FACTOR)
);
group_if!(prec, {
self.p(op)?;
self.unparse_expr(operand, prec)?;
})
}
ast::Expr::Lambda(ast::ExprLambda {
parameters,
body,
node_index: _,
range: _range,
}) => {
group_if!(precedence::TEST, {
if let Some(parameters) = parameters {
self.p("lambda ")?;
self.unparse_arguments(parameters)?;
} else {
self.p("lambda")?;
}
write!(self, ": {}", UnparseExpr::new(body, self.source))?;
})
}
ast::Expr::If(ast::ExprIf {
test,
body,
orelse,
node_index: _,
range: _range,
}) => {
group_if!(precedence::TEST, {
self.unparse_expr(body, precedence::TEST + 1)?;
self.p(" if ")?;
self.unparse_expr(test, precedence::TEST + 1)?;
self.p(" else ")?;
self.unparse_expr(orelse, precedence::TEST)?;
})
}
ast::Expr::Dict(ast::ExprDict {
items,
node_index: _,
range: _range,
}) => {
self.p("{")?;
let mut first = true;
for item in items {
self.p_delim(&mut first, ", ")?;
if let Some(k) = &item.key {
write!(self, "{}: ", UnparseExpr::new(k, self.source))?;
} else {
self.p("**")?;
}
self.unparse_expr(&item.value, level)?;
}
self.p("}")?;
}
ast::Expr::Set(ast::ExprSet {
elts,
node_index: _,
range: _range,
}) => {
self.p("{")?;
let mut first = true;
for v in elts {
self.p_delim(&mut first, ", ")?;
self.unparse_expr(v, precedence::TEST)?;
}
self.p("}")?;
}
ast::Expr::ListComp(ast::ExprListComp {
elt,
generators,
node_index: _,
range: _range,
}) => {
self.p("[")?;
self.unparse_expr(elt, precedence::TEST)?;
self.unparse_comp(generators)?;
self.p("]")?;
}
ast::Expr::SetComp(ast::ExprSetComp {
elt,
generators,
node_index: _,
range: _range,
}) => {
self.p("{")?;
self.unparse_expr(elt, precedence::TEST)?;
self.unparse_comp(generators)?;
self.p("}")?;
}
ast::Expr::DictComp(ast::ExprDictComp {
key,
value,
generators,
node_index: _,
range: _range,
}) => {
self.p("{")?;
self.unparse_expr(key, precedence::TEST)?;
self.p(": ")?;
self.unparse_expr(value, precedence::TEST)?;
self.unparse_comp(generators)?;
self.p("}")?;
}
ast::Expr::Generator(ast::ExprGenerator {
parenthesized: _,
elt,
generators,
node_index: _,
range: _range,
}) => {
self.p("(")?;
self.unparse_expr(elt, precedence::TEST)?;
self.unparse_comp(generators)?;
self.p(")")?;
}
ast::Expr::Await(ast::ExprAwait {
value,
node_index: _,
range: _range,
}) => {
group_if!(precedence::AWAIT, {
self.p("await ")?;
self.unparse_expr(value, precedence::ATOM)?;
})
}
ast::Expr::Yield(ast::ExprYield {
value,
node_index: _,
range: _range,
}) => {
if let Some(value) = value {
write!(self, "(yield {})", UnparseExpr::new(value, self.source))?;
} else {
self.p("(yield)")?;
}
}
ast::Expr::YieldFrom(ast::ExprYieldFrom {
value,
node_index: _,
range: _range,
}) => {
write!(
self,
"(yield from {})",
UnparseExpr::new(value, self.source)
)?;
}
ast::Expr::Compare(ast::ExprCompare {
left,
ops,
comparators,
node_index: _,
range: _range,
}) => {
group_if!(precedence::CMP, {
let new_lvl = precedence::CMP + 1;
self.unparse_expr(left, new_lvl)?;
for (op, cmp) in ops.iter().zip(comparators) {
self.p(" ")?;
self.p(op.as_str())?;
self.p(" ")?;
self.unparse_expr(cmp, new_lvl)?;
}
})
}
ast::Expr::Call(ast::ExprCall {
func,
arguments: ast::Arguments { args, keywords, .. },
node_index: _,
range: _range,
}) => {
self.unparse_expr(func, precedence::ATOM)?;
self.p("(")?;
if let (
[
ast::Expr::Generator(ast::ExprGenerator {
elt,
generators,
node_index: _,
range: _range,
..
}),
],
[],
) = (&**args, &**keywords)
{
// make sure a single genexpr doesn't get double parens
self.unparse_expr(elt, precedence::TEST)?;
self.unparse_comp(generators)?;
} else {
let mut first = true;
for arg in args {
self.p_delim(&mut first, ", ")?;
self.unparse_expr(arg, precedence::TEST)?;
}
for kw in keywords {
self.p_delim(&mut first, ", ")?;
if let Some(arg) = &kw.arg {
self.p_id(arg)?;
self.p("=")?;
} else {
self.p("**")?;
}
self.unparse_expr(&kw.value, precedence::TEST)?;
}
}
self.p(")")?;
}
ast::Expr::FString(ast::ExprFString { value, .. }) => self.unparse_fstring(value)?,
ast::Expr::TString(ast::ExprTString { value, .. }) => self.unparse_tstring(value)?,
ast::Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => {
if value.is_unicode() {
self.p("u")?
}
UnicodeEscape::new_repr(value.to_str().as_ref())
.str_repr()
.fmt(self.f)?
}
ast::Expr::BytesLiteral(ast::ExprBytesLiteral { value, .. }) => {
AsciiEscape::new_repr(&value.bytes().collect::