use crate::strip_python_comments;
use alloc::fmt;
use core::fmt::Display as _;
use ruff_python_ast as ast;
use ruff_text_size::{Ranged, TextSize, TextSlice};
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 {
const fn new(f: &'b mut fmt::Formatter 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 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::())
.bytes_repr()
.fmt(self.f)?
}
ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => match value {
ast::Number::Int(int) => int.fmt(self.f)?,
&ast::Number::Float(fp) => self.unparse_float(fp)?,
&ast::Number::Complex { real, imag } => self.unparse_complex(real, imag)?,
},
ast::Expr::Constant(ast::ExprConstant { value, .. }) => {
self.unparse_constant_value(value)?
}
ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => {
self.p(if *value { "True" } else { "False" })?
}
ast::Expr::NoneLiteral(ast::ExprNoneLiteral { .. }) => self.p("None")?,
ast::Expr::EllipsisLiteral(ast::ExprEllipsisLiteral { .. }) => self.p("...")?,
ast::Expr::Attribute(ast::ExprAttribute { value, attr, .. }) => {
self.unparse_expr(value, precedence::ATOM)?;
let period = if let ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
value: ast::Number::Int(_),
..
}) = value.as_ref()
{
" ."
} else {
"."
};
self.p(period)?;
self.p_id(attr)?;
}
ast::Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
self.unparse_expr(value, precedence::ATOM)?;
let lvl = precedence::TUPLE;
self.p("[")?;
self.unparse_expr(slice, lvl)?;
self.p("]")?;
}
ast::Expr::Starred(ast::ExprStarred { value, .. }) => {
self.p("*")?;
self.unparse_expr(value, precedence::EXPR)?;
}
ast::Expr::Name(ast::ExprName { id, .. }) => self.p(id.as_str())?,
ast::Expr::List(ast::ExprList { elts, .. }) => {
self.p("[")?;
let mut first = true;
for elt in elts {
self.p_delim(&mut first, ", ")?;
self.unparse_expr(elt, precedence::TEST)?;
}
self.p("]")?;
}
ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => {
if elts.is_empty() {
self.p("()")?;
} else {
group_if!(precedence::TUPLE, {
let mut first = true;
for elt in elts {
self.p_delim(&mut first, ", ")?;
self.unparse_expr(elt, precedence::TEST)?;
}
self.p_if(elts.len() == 1, ",")?;
})
}
}
ast::Expr::Slice(ast::ExprSlice {
lower,
upper,
step,
node_index: _,
range: _range,
..
}) => {
if let Some(lower) = lower {
self.unparse_expr(lower, precedence::TEST)?;
}
self.p(":")?;
if let Some(upper) = upper {
self.unparse_expr(upper, precedence::TEST)?;
}
if let Some(step) = step {
self.p(":")?;
self.unparse_expr(step, precedence::TEST)?;
}
}
ast::Expr::IpyEscapeCommand(_) => {}
}
Ok(())
}
fn unparse_arguments(&mut self, args: &ast::Parameters) -> fmt::Result {
let mut first = true;
for (i, arg) in args.posonlyargs.iter().chain(&args.args).enumerate() {
self.p_delim(&mut first, ", ")?;
self.unparse_function_arg(arg)?;
self.p_if(i + 1 == args.posonlyargs.len(), ", /")?;
}
if args.vararg.is_some() || !args.kwonlyargs.is_empty() {
self.p_delim(&mut first, ", ")?;
self.p("*")?;
}
if let Some(vararg) = &args.vararg {
self.unparse_arg(vararg)?;
}
for kwarg in &args.kwonlyargs {
self.p_delim(&mut first, ", ")?;
self.unparse_function_arg(kwarg)?;
}
if let Some(kwarg) = &args.kwarg {
self.p_delim(&mut first, ", ")?;
self.p("**")?;
self.unparse_arg(kwarg)?;
}
Ok(())
}
fn unparse_function_arg(&mut self, arg: &ast::ParameterWithDefault) -> fmt::Result {
self.unparse_arg(&arg.parameter)?;
if let Some(default) = &arg.default {
write!(self, "={}", UnparseExpr::new(default, self.source))?;
}
Ok(())
}
fn unparse_arg(&mut self, arg: &ast::Parameter) -> fmt::Result {
self.p_id(&arg.name)?;
if let Some(ann) = &arg.annotation {
write!(self, ": {}", UnparseExpr::new(ann, self.source))?;
}
Ok(())
}
fn unparse_comp(&mut self, generators: &[ast::Comprehension]) -> fmt::Result {
for comp in generators {
self.p(if comp.is_async {
" async for "
} else {
" for "
})?;
self.unparse_expr(&comp.target, precedence::TUPLE)?;
self.p(" in ")?;
self.unparse_expr(&comp.iter, precedence::TEST + 1)?;
for cond in &comp.ifs {
self.p(" if ")?;
self.unparse_expr(cond, precedence::TEST + 1)?;
}
}
Ok(())
}
fn unparse_fstring_body(&mut self, elements: &[ast::InterpolatedStringElement]) -> fmt::Result {
for elem in elements {
self.unparse_fstring_elem(elem)?;
}
Ok(())
}
fn unparse_formatted(
&mut self,
val: &ast::Expr,
debug_text: Option,
mut conversion: ast::ConversionFlag,
spec: Option,
) -> fmt::Result {
let buffered =
fmt::from_fn(|f| Unparser::new(f, self.source).unparse_expr(val, precedence::TEST + 1))
.to_string();
if let Some(debug_text) = debug_text {
let leading = debug_text.leading();
let trailing = debug_text.trailing();
self.p(leading)?;
self.p(self.source.source_text().slice(val.range()))?;
self.p(trailing)?;
if conversion == ast::ConversionFlag::None && spec.is_none() {
conversion = ast::ConversionFlag::Repr;
}
}
let brace = if buffered.starts_with('{') {
// put a space to avoid escaping the bracket
"{ "
} else {
// Preserve leading whitespace between '{' and the expression
let source_text = self.source.source_text();
let start = val.range().start().to_usize();
if start > 0
&& source_text
.as_bytes()
.get(start - 1)
.is_some_and(|b| b.is_ascii_whitespace())
{
"{ "
} else {
"{"
}
};
self.p(brace)?;
self.p(&buffered)?;
drop(buffered);
if conversion != ast::ConversionFlag::None {
self.p("!")?;
let buf = &[conversion as u8];
let c = core::str::from_utf8(buf).unwrap();
self.p(c)?;
}
if let Some(spec) = spec {
self.p(":")?;
self.unparse_fstring_body(&spec.elements)?;
}
self.p("}")?;
Ok(())
}
fn unparse_fstring_elem(&mut self, elem: &ast::InterpolatedStringElement) -> fmt::Result {
match elem {
ast::InterpolatedStringElement::Interpolation(ast::InterpolatedElement {
expression,
debug_text,
conversion,
format_spec,
..
}) => self.unparse_formatted(
expression,
debug_text.as_ref(),
*conversion,
format_spec.as_deref(),
),
ast::InterpolatedStringElement::Literal(ast::InterpolatedStringLiteralElement {
value,
..
}) => self.unparse_fstring_str(value),
}
}
fn unparse_fstring_str(&mut self, s: &str) -> fmt::Result {
let s = s.replace('{', "{{").replace('}', "}}");
self.p(&s)
}
fn unparse_fstring(&mut self, value: &ast::FStringValue) -> fmt::Result {
self.p("f")?;
let body = fmt::from_fn(|f| {
value.iter().try_for_each(|part| match part {
ast::FStringPartRef::Literal(lit) => f.write_str(lit),
ast::FStringPartRef::FString(ast::FString { elements, .. }) => {
Unparser::new(f, self.source).unparse_fstring_body(elements)
}
})
})
.to_string();
// .unparse_fstring_body(elements));
UnicodeEscape::new_repr(body.as_str().as_ref())
.str_repr()
.write(self.f)
}
fn unparse_tstring(&mut self, value: &ast::TStringValue) -> fmt::Result {
self.p("t")?;
let body = fmt::from_fn(|f| {
value.iter().try_for_each(|tstring| {
Unparser::new(f, self.source).unparse_tstring_body(&tstring.elements)
})
})
.to_string();
UnicodeEscape::new_repr(body.as_str().as_ref())
.str_repr()
.write(self.f)
}
fn unparse_tstring_body(&mut self, elements: &[ast::InterpolatedStringElement]) -> fmt::Result {
for element in elements {
match element {
ast::InterpolatedStringElement::Literal(literal) => {
self.unparse_fstring_str(literal)?;
}
ast::InterpolatedStringElement::Interpolation(interpolation) => {
self.unparse_tstring_interpolation(interpolation)?;
}
}
}
Ok(())
}
fn unparse_tstring_interpolation(
&mut self,
interpolation: &ast::InterpolatedElement,
) -> fmt::Result {
let source_conversion = interpolation.conversion;
let mut conversion = source_conversion;
let debug_parts = interpolation.debug_text.as_ref().map(|debug_text| {
(
strip_python_comments(debug_text.leading()),
strip_python_comments(
self.source
.source_text()
.slice(interpolation.expression.range()),
),
strip_python_comments(debug_text.trailing()),
)
});
if let Some((leading, source, trailing)) = &debug_parts {
self.p(leading)?;
self.p(source)?;
self.p(trailing)?;
if conversion == ast::ConversionFlag::None && interpolation.format_spec.is_none() {
conversion = ast::ConversionFlag::Repr;
}
}
let expression = if let Some(ast::ConstantValue::Str(value)) = &interpolation.runtime_str {
value.to_string()
} else if let Some((leading, source, trailing)) = &debug_parts {
let mut expression = leading.clone();
expression.push_str(source);
let equal = trailing
.rfind('=')
.expect("debug interpolation must contain '='");
expression.push_str(&trailing[..equal]);
expression.trim_end().to_owned()
} else {
let expression_range = interpolation.expression.range();
let after_brace = interpolation.range.start() + TextSize::new(1);
let mut expression_end = interpolation.format_spec.as_ref().map_or_else(
|| interpolation.range.end() - TextSize::new(1),
|format_spec| format_spec.range.start() - TextSize::new(1),
);
if source_conversion != ast::ConversionFlag::None {
expression_end -= TextSize::new(2);
}
if interpolation.range.start() < expression_range.start()
&& interpolation.range.end() >= expression_range.end()
&& after_brace UnparseExpr {
fn fmt(&self, f: &mut fmt::Formatter