From a2d0d398539d873cc9612d6cd8eb560cef7f55a6 Mon Sep 17 00:00:00 2001 From: Shahar Naveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Fri, 24 Oct 2025 18:44:48 +0300 Subject: [PATCH] Configurable "unparse mode" for `ruff_python_codegen::Generator` (#21041) Co-authored-by: Micha Reiser --- crates/ruff_python_codegen/src/generator.rs | 118 +++++++++++++------- crates/ruff_python_codegen/src/lib.rs | 2 +- 2 files changed, 77 insertions(+), 43 deletions(-) diff --git a/crates/ruff_python_codegen/src/generator.rs b/crates/ruff_python_codegen/src/generator.rs index af5b1dc80e..710e295f62 100644 --- a/crates/ruff_python_codegen/src/generator.rs +++ b/crates/ruff_python_codegen/src/generator.rs @@ -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, + /// 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) -> 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, + 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); diff --git a/crates/ruff_python_codegen/src/lib.rs b/crates/ruff_python_codegen/src/lib.rs index aeeab4a747..c30b35d932 100644 --- a/crates/ruff_python_codegen/src/lib.rs +++ b/crates/ruff_python_codegen/src/lib.rs @@ -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};