Preserve quotes in generated byte strings

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -313,6 +313,12 @@ impl<'a> Checker<'a> {
         ast::StringLiteralFlags::empty().with_quote_style(self.preferred_quote())
     }
 
+    /// Return the default bytestring flags a generated `ByteStringLiteral` node should use, given
+    /// where we are in the AST.
+    pub(crate) fn default_bytes_flags(&self) -> ast::BytesLiteralFlags {
+        ast::BytesLiteralFlags::empty().with_quote_style(self.preferred_quote())
+    }
+
     /// Returns the appropriate quoting for f-string by reversing the one used outside of
     /// the f-string.
     ///

crates/ruff_linter/src/rules/pyupgrade/rules/native_literals.rs

@@ -3,7 +3,7 @@ use std::str::FromStr;
 
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::{self as ast, Expr, Int, LiteralExpressionRef, StringLiteralFlags, UnaryOp};
+use ruff_python_ast::{self as ast, Expr, Int, LiteralExpressionRef, UnaryOp};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
@@ -33,15 +33,20 @@ impl FromStr for LiteralType {
 }
 
 impl LiteralType {
-    fn as_zero_value_expr(self, flags: StringLiteralFlags) -> Expr {
+    fn as_zero_value_expr(self, checker: &Checker) -> Expr {
         match self {
             LiteralType::Str => ast::StringLiteral {
                 value: Box::default(),
                 range: TextRange::default(),
-                flags,
+                flags: checker.default_string_flags(),
+            }
+            .into(),
+            LiteralType::Bytes => ast::BytesLiteral {
+                value: Box::default(),
+                range: TextRange::default(),
+                flags: checker.default_bytes_flags(),
             }
             .into(),
-            LiteralType::Bytes => ast::ExprBytesLiteral::default().into(),
             LiteralType::Int => ast::ExprNumberLiteral {
                 value: ast::Number::Int(Int::from(0u8)),
                 range: TextRange::default(),
@@ -191,7 +196,7 @@ pub(crate) fn native_literals(
                 return;
             }
 
-            let expr = literal_type.as_zero_value_expr(checker.default_string_flags());
+            let expr = literal_type.as_zero_value_expr(checker);
             let content = checker.generator().expr(&expr);
             diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
                 content,

crates/ruff_python_ast/src/nodes.rs

@@ -1634,14 +1634,14 @@ impl Debug for ConcatenatedStringLiteral {
 
 /// An AST node that represents either a single bytes literal or an implicitly
 /// concatenated bytes literals.
-#[derive(Clone, Debug, Default, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct ExprBytesLiteral {
     pub range: TextRange,
     pub value: BytesLiteralValue,
 }
 
 /// The value representing a [`ExprBytesLiteral`].
-#[derive(Clone, Debug, Default, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct BytesLiteralValue {
     inner: BytesLiteralValueInner,
 }
@@ -1714,6 +1714,16 @@ impl BytesLiteralValue {
     pub fn bytes(&self) -> impl Iterator<Item = u8> + '_ {
         self.iter().flat_map(|part| part.as_slice().iter().copied())
     }
+
+    /// Returns the [`BytesLiteralFlags`] associated with this literal.
+    ///
+    /// For an implicitly concatenated literal, it returns the flags for the first literal.
+    pub fn flags(&self) -> BytesLiteralFlags {
+        self.iter()
+            .next()
+            .expect("There should always be at least one literal in an `ExprBytesLiteral` node")
+            .flags
+    }
 }
 
 impl<'a> IntoIterator for &'a BytesLiteralValue {
@@ -1771,12 +1781,6 @@ enum BytesLiteralValueInner {
     Concatenated(Vec<BytesLiteral>),
 }
 
-impl Default for BytesLiteralValueInner {
-    fn default() -> Self {
-        Self::Single(BytesLiteral::default())
-    }
-}
-
 bitflags! {
     #[derive(Default, Copy, Clone, PartialEq, Eq, Hash)]
     struct BytesLiteralFlagsInner: u8 {
@@ -1805,10 +1809,33 @@ bitflags! {
 
 /// Flags that can be queried to obtain information
 /// regarding the prefixes and quotes used for a bytes literal.
-#[derive(Default, Copy, Clone, Eq, PartialEq, Hash)]
+///
+/// ## Notes on usage
+///
+/// If you're using a `Generator` from the `ruff_python_codegen` crate to generate a lint-rule fix
+/// from an existing bytes literal, consider passing along the [`BytesLiteral::flags`] field or the
+/// result of the [`BytesLiteralValue::flags`] method. If you don't have an existing literal but
+/// have a `Checker` from the `ruff_linter` crate available, consider using
+/// `Checker::default_bytes_flags` to create instances of this struct; this method will properly
+/// handle surrounding f-strings. For usage that doesn't fit into one of these categories, the
+/// public constructor [`BytesLiteralFlags::empty`] can be used.
+#[derive(Copy, Clone, Eq, PartialEq, Hash)]
 pub struct BytesLiteralFlags(BytesLiteralFlagsInner);
 
 impl BytesLiteralFlags {
+    /// Construct a new [`BytesLiteralFlags`] with **no flags set**.
+    ///
+    /// See [`BytesLiteralFlags::with_quote_style`], [`BytesLiteralFlags::with_triple_quotes`], and
+    /// [`BytesLiteralFlags::with_prefix`] for ways of setting the quote style (single or double),
+    /// enabling triple quotes, and adding prefixes (such as `r`), respectively.
+    ///
+    /// See the documentation for [`BytesLiteralFlags`] for additional caveats on this constructor,
+    /// and situations in which alternative ways to construct this struct should be used, especially
+    /// when writing lint rules.
+    pub fn empty() -> Self {
+        Self(BytesLiteralFlagsInner::empty())
+    }
+
     #[must_use]
     pub fn with_quote_style(mut self, quote_style: Quote) -> Self {
         self.0
@@ -1894,7 +1921,7 @@ impl fmt::Debug for BytesLiteralFlags {
 
 /// An AST node that represents a single bytes literal which is part of an
 /// [`ExprBytesLiteral`].
-#[derive(Clone, Debug, Default, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct BytesLiteral {
     pub range: TextRange,
     pub value: Box<[u8]>,
@@ -1920,7 +1947,7 @@ impl BytesLiteral {
         Self {
             range,
             value: Box::new([]),
-            flags: BytesLiteralFlags::default().with_invalid(),
+            flags: BytesLiteralFlags::empty().with_invalid(),
         }
     }
 }
@@ -2173,7 +2200,7 @@ impl From<AnyStringFlags> for BytesLiteralFlags {
                 value.prefix()
             )
         };
-        let new = BytesLiteralFlags::default()
+        let new = BytesLiteralFlags::empty()
             .with_quote_style(value.quote_style())
             .with_prefix(bytestring_prefix);
         if value.is_triple_quoted() {

crates/ruff_python_codegen/src/generator.rs

@@ -153,8 +153,8 @@ impl<'a> Generator<'a> {
         self.p(s.as_str());
     }
 
-    fn p_bytes_repr(&mut self, s: &[u8]) {
-        let escape = AsciiEscape::with_preferred_quote(s, self.quote);
+    fn p_bytes_repr(&mut self, s: &[u8], quote: Quote) {
+        let escape = AsciiEscape::with_preferred_quote(s, quote);
         if let Some(len) = escape.layout().len {
             self.buffer.reserve(len);
         }
@@ -1100,7 +1100,7 @@ impl<'a> Generator<'a> {
                 let mut first = true;
                 for bytes_literal in value {
                     self.p_delim(&mut first, " ");
-                    self.p_bytes_repr(&bytes_literal.value);
+                    self.p_bytes_repr(&bytes_literal.value, bytes_literal.flags.quote_style());
                 }
             }
             Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => {
@@ -1739,11 +1739,11 @@ class Foo:
 
     #[test]
     fn quote() {
-        assert_eq!(round_trip(r#""hello""#), r#""hello""#);
+        assert_round_trip!(r#""hello""#);
         assert_round_trip!(r"'hello'");
         assert_round_trip!(r"u'hello'");
         assert_round_trip!(r"r'hello'");
-        assert_eq!(round_trip(r"b'hello'"), r#"b"hello""#);
+        assert_round_trip!(r"b'hello'");
         assert_eq!(round_trip(r#"("abc" "def" "ghi")"#), r#""abc" "def" "ghi""#);
         assert_eq!(round_trip(r#""he\"llo""#), r#"'he"llo'"#);
         assert_eq!(round_trip(r#"f"abc{'def'}{1}""#), r#"f"abc{'def'}{1}""#);
@@ -1806,25 +1806,28 @@ if True:
                     $end,
                 );
             };
+            ($quote:expr, $start:expr) => {
+                round_trip_with!($quote, $start, $start);
+            };
         }
 
-        // setting Generator::quote works for bytestrings
-        round_trip_with!(Quote::Double, r#"b"hello""#, r#"b"hello""#);
-        round_trip_with!(Quote::Single, r#"b"hello""#, r"b'hello'");
-        round_trip_with!(Quote::Double, r"b'hello'", r#"b"hello""#);
-        round_trip_with!(Quote::Single, r"b'hello'", r"b'hello'");
-
-        // and for f-strings
+        // setting Generator::quote works for f-strings
         round_trip_with!(Quote::Double, r#"f"hello""#, r#"f"hello""#);
         round_trip_with!(Quote::Single, r#"f"hello""#, r"f'hello'");
         round_trip_with!(Quote::Double, r"f'hello'", r#"f"hello""#);
         round_trip_with!(Quote::Single, r"f'hello'", r"f'hello'");
 
-        // but not for string literals, where the `Quote` is taken directly from their flags
-        round_trip_with!(Quote::Double, r#""hello""#, r#""hello""#);
-        round_trip_with!(Quote::Single, r#""hello""#, r#""hello""#); // no effect
-        round_trip_with!(Quote::Double, r"'hello'", r#"'hello'"#); // no effect
-        round_trip_with!(Quote::Single, r"'hello'", r"'hello'");
+        // but not for bytestrings
+        round_trip_with!(Quote::Double, r#"b"hello""#);
+        round_trip_with!(Quote::Single, r#"b"hello""#); // no effect
+        round_trip_with!(Quote::Double, r"b'hello'"); // no effect
+        round_trip_with!(Quote::Single, r"b'hello'");
+
+        // or for string literals, where the `Quote` is taken directly from their flags
+        round_trip_with!(Quote::Double, r#""hello""#);
+        round_trip_with!(Quote::Single, r#""hello""#); // no effect
+        round_trip_with!(Quote::Double, r"'hello'"); // no effect
+        round_trip_with!(Quote::Single, r"'hello'");
     }
 
     #[test]

crates/ruff_python_formatter/tests/normalizer.rs

@@ -97,7 +97,7 @@ impl Transformer for Normalizer {
                         bytes.value = ast::BytesLiteralValue::single(ast::BytesLiteral {
                             value: bytes.value.bytes().collect(),
                             range: bytes.range,
-                            flags: BytesLiteralFlags::default(),
+                            flags: BytesLiteralFlags::empty(),
                         });
                     }
                 }