mirror of
https://github.com/astral-sh/ruff.git
synced 2025-11-17 19:27:11 +00:00
Configurable "unparse mode" for ruff_python_codegen::Generator (#21041)
Co-authored-by: Micha Reiser <micha@reiser.io>
This commit is contained in:
parent
f36fa7d6c1
commit
a2d0d39853
2 changed files with 77 additions and 43 deletions
|
|
@ -17,6 +17,7 @@ use ruff_source_file::LineEnding;
|
|||
use super::stylist::{Indentation, Stylist};
|
||||
|
||||
mod precedence {
|
||||
pub(crate) const MIN: u8 = 0;
|
||||
pub(crate) const NAMED_EXPR: u8 = 1;
|
||||
pub(crate) const ASSIGN: u8 = 3;
|
||||
pub(crate) const ANN_ASSIGN: u8 = 5;
|
||||
|
|
@ -63,13 +64,36 @@ mod precedence {
|
|||
pub(crate) const MAX: u8 = 63;
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub enum Mode {
|
||||
/// Ruff's default unparsing behaviour.
|
||||
#[default]
|
||||
Default,
|
||||
/// Emits same output as [`ast.unparse`](https://docs.python.org/3/library/ast.html#ast.unparse).
|
||||
AstUnparse,
|
||||
}
|
||||
|
||||
impl Mode {
|
||||
/// Quote style to use.
|
||||
///
|
||||
/// - [`Default`](`Mode::Default`): Output of `[AnyStringFlags.quote_style`].
|
||||
/// - [`AstUnparse`](`Mode::AstUnparse`): Always return [`Quote::Single`].
|
||||
#[must_use]
|
||||
fn quote_style(&self, flags: impl StringFlags) -> Quote {
|
||||
match self {
|
||||
Self::Default => flags.quote_style(),
|
||||
Self::AstUnparse => Quote::Single,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Generator<'a> {
|
||||
/// The indentation style to use.
|
||||
indent: &'a Indentation,
|
||||
/// The line ending to use.
|
||||
line_ending: LineEnding,
|
||||
/// Preferred quote style to use. For more info see [`Generator::with_preferred_quote`].
|
||||
preferred_quote: Option<Quote>,
|
||||
/// Unparsed code style. See [`Mode`] for more info.
|
||||
mode: Mode,
|
||||
buffer: String,
|
||||
indent_depth: usize,
|
||||
num_newlines: usize,
|
||||
|
|
@ -81,7 +105,7 @@ impl<'a> From<&'a Stylist<'a>> for Generator<'a> {
|
|||
Self {
|
||||
indent: stylist.indentation(),
|
||||
line_ending: stylist.line_ending(),
|
||||
preferred_quote: None,
|
||||
mode: Mode::default(),
|
||||
buffer: String::new(),
|
||||
indent_depth: 0,
|
||||
num_newlines: 0,
|
||||
|
|
@ -96,7 +120,7 @@ impl<'a> Generator<'a> {
|
|||
// Style preferences.
|
||||
indent,
|
||||
line_ending,
|
||||
preferred_quote: None,
|
||||
mode: Mode::Default,
|
||||
// Internal state.
|
||||
buffer: String::new(),
|
||||
indent_depth: 0,
|
||||
|
|
@ -105,13 +129,10 @@ impl<'a> Generator<'a> {
|
|||
}
|
||||
}
|
||||
|
||||
/// Set a preferred quote style for generated source code.
|
||||
///
|
||||
/// - If [`None`], the generator will attempt to preserve the existing quote style whenever possible.
|
||||
/// - If [`Some`], the generator will prefer the specified quote style, ignoring the one found in the source.
|
||||
/// Sets the mode for code unparsing.
|
||||
#[must_use]
|
||||
pub fn with_preferred_quote(mut self, quote: Option<Quote>) -> Self {
|
||||
self.preferred_quote = quote;
|
||||
pub fn with_mode(mut self, mode: Mode) -> Self {
|
||||
self.mode = mode;
|
||||
self
|
||||
}
|
||||
|
||||
|
|
@ -173,7 +194,7 @@ impl<'a> Generator<'a> {
|
|||
return;
|
||||
}
|
||||
}
|
||||
let quote_style = self.preferred_quote.unwrap_or_else(|| flags.quote_style());
|
||||
let quote_style = self.mode.quote_style(flags);
|
||||
let escape = AsciiEscape::with_preferred_quote(s, quote_style);
|
||||
if let Some(len) = escape.layout().len {
|
||||
self.buffer.reserve(len);
|
||||
|
|
@ -193,7 +214,7 @@ impl<'a> Generator<'a> {
|
|||
}
|
||||
self.p(flags.prefix().as_str());
|
||||
|
||||
let quote_style = self.preferred_quote.unwrap_or_else(|| flags.quote_style());
|
||||
let quote_style = self.mode.quote_style(flags);
|
||||
let escape = UnicodeEscape::with_preferred_quote(s, quote_style);
|
||||
if let Some(len) = escape.layout().len {
|
||||
self.buffer.reserve(len);
|
||||
|
|
@ -1303,7 +1324,11 @@ impl<'a> Generator<'a> {
|
|||
if tuple.is_empty() {
|
||||
self.p("()");
|
||||
} else {
|
||||
group_if!(precedence::TUPLE, {
|
||||
let lvl = match self.mode {
|
||||
Mode::Default => precedence::TUPLE,
|
||||
Mode::AstUnparse => precedence::MIN,
|
||||
};
|
||||
group_if!(lvl, {
|
||||
let mut first = true;
|
||||
for item in tuple {
|
||||
self.p_delim(&mut first, ", ");
|
||||
|
|
@ -1525,7 +1550,7 @@ impl<'a> Generator<'a> {
|
|||
return;
|
||||
}
|
||||
|
||||
let quote_style = self.preferred_quote.unwrap_or_else(|| flags.quote_style());
|
||||
let quote_style = self.mode.quote_style(flags);
|
||||
let escape = UnicodeEscape::with_preferred_quote(&s, quote_style);
|
||||
if let Some(len) = escape.layout().len {
|
||||
self.buffer.reserve(len);
|
||||
|
|
@ -1552,8 +1577,8 @@ impl<'a> Generator<'a> {
|
|||
) {
|
||||
self.p(flags.prefix().as_str());
|
||||
|
||||
let flags =
|
||||
flags.with_quote_style(self.preferred_quote.unwrap_or_else(|| flags.quote_style()));
|
||||
let quote_style = self.mode.quote_style(flags);
|
||||
let flags = flags.with_quote_style(quote_style);
|
||||
self.p(flags.quote_str());
|
||||
self.unparse_interpolated_string_body(values, flags);
|
||||
self.p(flags.quote_str());
|
||||
|
|
@ -1586,14 +1611,13 @@ impl<'a> Generator<'a> {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use ruff_python_ast::str::Quote;
|
||||
use ruff_python_ast::{Mod, ModModule};
|
||||
use ruff_python_parser::{self, Mode, ParseOptions, parse_module};
|
||||
use ruff_source_file::LineEnding;
|
||||
|
||||
use crate::stylist::Indentation;
|
||||
|
||||
use super::Generator;
|
||||
use super::{Generator, Mode as UnparseMode};
|
||||
|
||||
fn round_trip(contents: &str) -> String {
|
||||
let indentation = Indentation::default();
|
||||
|
|
@ -1605,16 +1629,15 @@ mod tests {
|
|||
}
|
||||
|
||||
/// Like [`round_trip`] but configure the [`Generator`] with the requested
|
||||
/// `indentation`, `line_ending` and `preferred_quote` settings.
|
||||
/// `indentation`, `line_ending` and `unparse_mode` settings.
|
||||
fn round_trip_with(
|
||||
indentation: &Indentation,
|
||||
line_ending: LineEnding,
|
||||
preferred_quote: Option<Quote>,
|
||||
unparse_mode: UnparseMode,
|
||||
contents: &str,
|
||||
) -> String {
|
||||
let module = parse_module(contents).unwrap();
|
||||
let mut generator =
|
||||
Generator::new(indentation, line_ending).with_preferred_quote(preferred_quote);
|
||||
let mut generator = Generator::new(indentation, line_ending).with_mode(unparse_mode);
|
||||
generator.unparse_suite(module.suite());
|
||||
generator.generate()
|
||||
}
|
||||
|
|
@ -1814,6 +1837,7 @@ except* Exception as e:
|
|||
type Y = str"
|
||||
);
|
||||
assert_eq!(round_trip(r"x = (1, 2, 3)"), r"x = 1, 2, 3");
|
||||
assert_eq!(round_trip(r"x = (1, (2, 3))"), r"x = 1, (2, 3)");
|
||||
assert_eq!(round_trip(r"-(1) + ~(2) + +(3)"), r"-1 + ~2 + +3");
|
||||
assert_round_trip!(
|
||||
r"def f():
|
||||
|
|
@ -2000,7 +2024,7 @@ if True:
|
|||
round_trip_with(
|
||||
&Indentation::new(" ".to_string()),
|
||||
LineEnding::default(),
|
||||
None,
|
||||
UnparseMode::Default,
|
||||
r"
|
||||
if True:
|
||||
pass
|
||||
|
|
@ -2018,7 +2042,7 @@ if True:
|
|||
round_trip_with(
|
||||
&Indentation::new(" ".to_string()),
|
||||
LineEnding::default(),
|
||||
None,
|
||||
UnparseMode::Default,
|
||||
r"
|
||||
if True:
|
||||
pass
|
||||
|
|
@ -2036,7 +2060,7 @@ if True:
|
|||
round_trip_with(
|
||||
&Indentation::new("\t".to_string()),
|
||||
LineEnding::default(),
|
||||
None,
|
||||
UnparseMode::Default,
|
||||
r"
|
||||
if True:
|
||||
pass
|
||||
|
|
@ -2058,7 +2082,7 @@ if True:
|
|||
round_trip_with(
|
||||
&Indentation::default(),
|
||||
LineEnding::Lf,
|
||||
None,
|
||||
UnparseMode::Default,
|
||||
"if True:\n print(42)",
|
||||
),
|
||||
"if True:\n print(42)",
|
||||
|
|
@ -2068,7 +2092,7 @@ if True:
|
|||
round_trip_with(
|
||||
&Indentation::default(),
|
||||
LineEnding::CrLf,
|
||||
None,
|
||||
UnparseMode::Default,
|
||||
"if True:\n print(42)",
|
||||
),
|
||||
"if True:\r\n print(42)",
|
||||
|
|
@ -2078,30 +2102,40 @@ if True:
|
|||
round_trip_with(
|
||||
&Indentation::default(),
|
||||
LineEnding::Cr,
|
||||
None,
|
||||
UnparseMode::Default,
|
||||
"if True:\n print(42)",
|
||||
),
|
||||
"if True:\r print(42)",
|
||||
);
|
||||
}
|
||||
|
||||
#[test_case::test_case(r#""'hello'""#, r#""'hello'""#, Quote::Single ; "basic str ignored")]
|
||||
#[test_case::test_case(r#"b"'hello'""#, r#"b"'hello'""#, Quote::Single ; "basic bytes ignored")]
|
||||
#[test_case::test_case("'hello'", r#""hello""#, Quote::Double ; "basic str double")]
|
||||
#[test_case::test_case(r#""hello""#, "'hello'", Quote::Single ; "basic str single")]
|
||||
#[test_case::test_case("b'hello'", r#"b"hello""#, Quote::Double ; "basic bytes double")]
|
||||
#[test_case::test_case(r#"b"hello""#, "b'hello'", Quote::Single ; "basic bytes single")]
|
||||
#[test_case::test_case(r#""hello""#, r#""hello""#, Quote::Double ; "remain str double")]
|
||||
#[test_case::test_case("'hello'", "'hello'", Quote::Single ; "remain str single")]
|
||||
#[test_case::test_case("x: list['str']", r#"x: list["str"]"#, Quote::Double ; "type ann double")]
|
||||
#[test_case::test_case(r#"x: list["str"]"#, "x: list['str']", Quote::Single ; "type ann single")]
|
||||
#[test_case::test_case("f'hello'", r#"f"hello""#, Quote::Double ; "basic fstring double")]
|
||||
#[test_case::test_case(r#"f"hello""#, "f'hello'", Quote::Single ; "basic fstring single")]
|
||||
fn preferred_quote(inp: &str, out: &str, quote: Quote) {
|
||||
#[test_case::test_case(r#""'hello'""#, r#""'hello'""# ; "basic str ignored")]
|
||||
#[test_case::test_case(r#"b"'hello'""#, r#"b"'hello'""# ; "basic bytes ignored")]
|
||||
#[test_case::test_case(r#""hello""#, "'hello'" ; "basic str single")]
|
||||
#[test_case::test_case(r#"b"hello""#, "b'hello'" ; "basic bytes single")]
|
||||
#[test_case::test_case("'hello'", "'hello'" ; "remain str single")]
|
||||
#[test_case::test_case(r#"x: list["str"]"#, "x: list['str']" ; "type ann single")]
|
||||
#[test_case::test_case(r#"f"hello""#, "f'hello'" ; "basic fstring single")]
|
||||
fn ast_unparse_quote(inp: &str, out: &str) {
|
||||
let got = round_trip_with(
|
||||
&Indentation::default(),
|
||||
LineEnding::default(),
|
||||
Some(quote),
|
||||
UnparseMode::AstUnparse,
|
||||
inp,
|
||||
);
|
||||
assert_eq!(got, out);
|
||||
}
|
||||
|
||||
#[test_case::test_case("a,", "(a,)" ; "basic single")]
|
||||
#[test_case::test_case("a, b", "(a, b)" ; "basic multi")]
|
||||
#[test_case::test_case("x = a,", "x = (a,)" ; "basic assign single")]
|
||||
#[test_case::test_case("x = a, b", "x = (a, b)" ; "basic assign multi")]
|
||||
#[test_case::test_case("a, (b, c)", "(a, (b, c))" ; "nested")]
|
||||
fn ast_tuple_parentheses(inp: &str, out: &str) {
|
||||
let got = round_trip_with(
|
||||
&Indentation::default(),
|
||||
LineEnding::default(),
|
||||
UnparseMode::AstUnparse,
|
||||
inp,
|
||||
);
|
||||
assert_eq!(got, out);
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
pub use generator::Generator;
|
||||
pub use generator::{Generator, Mode};
|
||||
use ruff_python_parser::{ParseError, parse_module};
|
||||
pub use stylist::{Indentation, Stylist};
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue