diff --git a/crates/ruff_python_codegen/src/generator.rs b/crates/ruff_python_codegen/src/generator.rs index 1bd307da8d..af5b1dc80e 100644 --- a/crates/ruff_python_codegen/src/generator.rs +++ b/crates/ruff_python_codegen/src/generator.rs @@ -3,6 +3,7 @@ use std::fmt::Write; use std::ops::Deref; +use ruff_python_ast::str::Quote; use ruff_python_ast::{ self as ast, Alias, AnyStringFlags, ArgOrKeyword, BoolOp, BytesLiteralFlags, CmpOp, Comprehension, ConversionFlag, DebugText, ExceptHandler, Expr, Identifier, MatchCase, Operator, @@ -67,6 +68,8 @@ pub struct Generator<'a> { 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, buffer: String, indent_depth: usize, num_newlines: usize, @@ -78,6 +81,7 @@ impl<'a> From<&'a Stylist<'a>> for Generator<'a> { Self { indent: stylist.indentation(), line_ending: stylist.line_ending(), + preferred_quote: None, buffer: String::new(), indent_depth: 0, num_newlines: 0, @@ -92,6 +96,7 @@ impl<'a> Generator<'a> { // Style preferences. indent, line_ending, + preferred_quote: None, // Internal state. buffer: String::new(), indent_depth: 0, @@ -100,6 +105,16 @@ 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. + #[must_use] + pub fn with_preferred_quote(mut self, quote: Option) -> Self { + self.preferred_quote = quote; + self + } + /// Generate source code from a [`Stmt`]. pub fn stmt(mut self, stmt: &Stmt) -> String { self.unparse_stmt(stmt); @@ -158,7 +173,8 @@ impl<'a> Generator<'a> { return; } } - let escape = AsciiEscape::with_preferred_quote(s, flags.quote_style()); + let quote_style = self.preferred_quote.unwrap_or_else(|| flags.quote_style()); + let escape = AsciiEscape::with_preferred_quote(s, quote_style); if let Some(len) = escape.layout().len { self.buffer.reserve(len); } @@ -176,7 +192,9 @@ impl<'a> Generator<'a> { return; } self.p(flags.prefix().as_str()); - let escape = UnicodeEscape::with_preferred_quote(s, flags.quote_style()); + + let quote_style = self.preferred_quote.unwrap_or_else(|| flags.quote_style()); + let escape = UnicodeEscape::with_preferred_quote(s, quote_style); if let Some(len) = escape.layout().len { self.buffer.reserve(len); } @@ -1506,7 +1524,9 @@ impl<'a> Generator<'a> { self.buffer += &s; return; } - let escape = UnicodeEscape::with_preferred_quote(&s, flags.quote_style()); + + let quote_style = self.preferred_quote.unwrap_or_else(|| flags.quote_style()); + let escape = UnicodeEscape::with_preferred_quote(&s, quote_style); if let Some(len) = escape.layout().len { self.buffer.reserve(len); } @@ -1531,6 +1551,9 @@ impl<'a> Generator<'a> { flags: AnyStringFlags, ) { self.p(flags.prefix().as_str()); + + let flags = + flags.with_quote_style(self.preferred_quote.unwrap_or_else(|| flags.quote_style())); self.p(flags.quote_str()); self.unparse_interpolated_string_body(values, flags); self.p(flags.quote_str()); @@ -1563,6 +1586,7 @@ 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; @@ -1580,15 +1604,17 @@ mod tests { generator.generate() } - /// Like [`round_trip`] but configure the [`Generator`] with the requested `indentation` and - /// `line_ending` settings. + /// Like [`round_trip`] but configure the [`Generator`] with the requested + /// `indentation`, `line_ending` and `preferred_quote` settings. fn round_trip_with( indentation: &Indentation, line_ending: LineEnding, + preferred_quote: Option, contents: &str, ) -> String { let module = parse_module(contents).unwrap(); - let mut generator = Generator::new(indentation, line_ending); + let mut generator = + Generator::new(indentation, line_ending).with_preferred_quote(preferred_quote); generator.unparse_suite(module.suite()); generator.generate() } @@ -1974,6 +2000,7 @@ if True: round_trip_with( &Indentation::new(" ".to_string()), LineEnding::default(), + None, r" if True: pass @@ -1991,6 +2018,7 @@ if True: round_trip_with( &Indentation::new(" ".to_string()), LineEnding::default(), + None, r" if True: pass @@ -2008,6 +2036,7 @@ if True: round_trip_with( &Indentation::new("\t".to_string()), LineEnding::default(), + None, r" if True: pass @@ -2029,6 +2058,7 @@ if True: round_trip_with( &Indentation::default(), LineEnding::Lf, + None, "if True:\n print(42)", ), "if True:\n print(42)", @@ -2038,6 +2068,7 @@ if True: round_trip_with( &Indentation::default(), LineEnding::CrLf, + None, "if True:\n print(42)", ), "if True:\r\n print(42)", @@ -2047,9 +2078,32 @@ if True: round_trip_with( &Indentation::default(), LineEnding::Cr, + None, "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) { + let got = round_trip_with( + &Indentation::default(), + LineEnding::default(), + Some(quote), + inp, + ); + assert_eq!(got, out); + } }