Add support for reformatting byte strings (#3176)

crates/ruff_python_formatter/src/core/helpers.rs

@@ -1,11 +1,12 @@
-use ruff_python::string::{
-    SINGLE_QUOTE_PREFIXES, SINGLE_QUOTE_SUFFIXES, TRIPLE_QUOTE_PREFIXES, TRIPLE_QUOTE_SUFFIXES,
-};
-
-/// Return the leading quote string for a docstring (e.g., `"""`).
+/// Return the leading quote for a string or byte literal (e.g., `"""`).
 pub fn leading_quote(content: &str) -> Option<&str> {
     if let Some(first_line) = content.lines().next() {
-        for pattern in TRIPLE_QUOTE_PREFIXES.iter().chain(SINGLE_QUOTE_PREFIXES) {
+        for pattern in ruff_python::str::TRIPLE_QUOTE_PREFIXES
+            .iter()
+            .chain(ruff_python::bytes::TRIPLE_QUOTE_PREFIXES)
+            .chain(ruff_python::str::SINGLE_QUOTE_PREFIXES)
+            .chain(ruff_python::bytes::SINGLE_QUOTE_PREFIXES)
+        {
             if first_line.starts_with(pattern) {
                 return Some(pattern);
             }
@@ -14,10 +15,38 @@ pub fn leading_quote(content: &str) -> Option<&str> {
     None
 }
 
-/// Return the trailing quote string for a docstring (e.g., `"""`).
+/// Return the trailing quote string for a string or byte literal (e.g., `"""`).
 pub fn trailing_quote(content: &str) -> Option<&&str> {
-    TRIPLE_QUOTE_SUFFIXES
+    ruff_python::str::TRIPLE_QUOTE_SUFFIXES
         .iter()
-        .chain(SINGLE_QUOTE_SUFFIXES)
+        .chain(ruff_python::str::SINGLE_QUOTE_SUFFIXES)
         .find(|&pattern| content.ends_with(pattern))
 }
+
+#[cfg(test)]
+mod tests {
+
+    #[test]
+    fn test_prefixes() {
+        let prefixes = ruff_python::str::TRIPLE_QUOTE_PREFIXES
+            .iter()
+            .chain(ruff_python::bytes::TRIPLE_QUOTE_PREFIXES)
+            .chain(ruff_python::str::SINGLE_QUOTE_PREFIXES)
+            .chain(ruff_python::bytes::SINGLE_QUOTE_PREFIXES)
+            .collect::<Vec<_>>();
+        for i in 1..prefixes.len() {
+            for j in 0..i - 1 {
+                if i != j {
+                    if prefixes[i].starts_with(prefixes[j]) {
+                        assert!(
+                            !prefixes[i].starts_with(prefixes[j]),
+                            "Prefixes are not unique: {} starts with {}",
+                            prefixes[i],
+                            prefixes[j]
+                        );
+                    }
+                }
+            }
+        }
+    }
+}

crates/ruff_python_formatter/src/format/expr.rs

@@ -652,7 +652,7 @@ fn format_constant(
                 write!(f, [text("False")])?;
             }
         }
-        Constant::Str(_) => write!(f, [string_literal(expr)])?,
+        Constant::Str(_) | Constant::Bytes(_) => write!(f, [string_literal(expr)])?,
         _ => write!(f, [literal(Range::from_located(expr))])?,
     }
     Ok(())

crates/ruff_python_formatter/src/format/strings.rs

@@ -37,6 +37,7 @@ impl Format<ASTFormatContext<'_>> for StringLiteralPart {
             }
         }
 
+        // Retain raw prefixes.
         let mut is_raw = false;
         if leading_quote.contains('r') {
             is_raw = true;
@@ -46,6 +47,11 @@ impl Format<ASTFormatContext<'_>> for StringLiteralPart {
             f.write_element(FormatElement::StaticText { text: "R" })?;
         }
 
+        // Normalize bytes literals to use b"...".
+        if leading_quote.contains('b') || leading_quote.contains('B') {
+            f.write_element(FormatElement::StaticText { text: "b" })?;
+        }
+
         if trailing_quote.len() == 1 {
             // Single-quoted string.
             if dquotes == 0 || squotes > 0 {

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap

@@ -267,14 +267,11 @@ last_call()
 ```diff
 --- Black
 +++ Ruff
-@@ -1,5 +1,6 @@
+@@ -1,3 +1,4 @@
 +...
  "some_string"
--b"\\xa3"
-+b'\\xa3'
+ b"\\xa3"
  Name
- None
- True
 @@ -38,7 +39,8 @@
  lambda a, b, c=True, *, d=(1 << v2), e="str": a
  lambda a, b, c=True, *vararg, d=(v1 << 2), e="str", **kwargs: a + b
@@ -285,15 +282,6 @@ last_call()
      "port1": port1_resource,
      "port2": port2_resource,
  }[port_id]
-@@ -56,7 +58,7 @@
- {"2.7": dead, "3.7": (long_live or die_hard), **{"3.6": verygood}}
- {**a, **b, **c}
- {"2.7", "3.6", "3.7", "3.8", "3.9", ("4.0" if gilectomy else "3.10")}
--({"a": "b"}, (True or False), (+value), "string", b"bytes") or None
-+({"a": "b"}, (True or False), (+value), "string", b'bytes') or None
- ()
- (1,)
- (1, 2)
 @@ -100,7 +102,8 @@
  {a: b * -2 for a, b in dictionary.items()}
  {
@@ -470,7 +458,7 @@ last_call()
 ```py
 ...
 "some_string"
-b'\\xa3'
+b"\\xa3"
 Name
 None
 True
@@ -528,7 +516,7 @@ str or None if (1 if True else 2) else str or bytes or None
 {"2.7": dead, "3.7": (long_live or die_hard), **{"3.6": verygood}}
 {**a, **b, **c}
 {"2.7", "3.6", "3.7", "3.8", "3.9", ("4.0" if gilectomy else "3.10")}
-({"a": "b"}, (True or False), (+value), "string", b'bytes') or None
+({"a": "b"}, (True or False), (+value), "string", b"bytes") or None
 ()
 (1,)
 (1, 2)

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__string_prefixes_py.snap

@@ -33,24 +33,21 @@ def docstring_multiline():
 ```diff
 --- Black
 +++ Ruff
-@@ -1,13 +1,13 @@
+@@ -1,12 +1,12 @@
  #!/usr/bin/env python3
  
  name = "Łukasz"
 -(f"hello {name}", f"hello {name}")
--(b"", b"")
 +(f"hello {name}", F"hello {name}")
-+(b"", B"")
+ (b"", b"")
  ("", "")
  (r"", R"")
  
 -(rf"", rf"", Rf"", Rf"", rf"", rf"", Rf"", Rf"")
--(rb"", rb"", Rb"", Rb"", rb"", rb"", Rb"", Rb"")
 +(rf"", fr"", Rf"", fR"", rF"", Fr"", RF"", FR"")
-+(rb"", br"", Rb"", bR"", rB"", Br"", RB"", BR"")
+ (rb"", rb"", Rb"", Rb"", rb"", rb"", Rb"", Rb"")
  
  
- def docstring_singleline():
 ```
 
 ## Ruff Output
@@ -60,12 +57,12 @@ def docstring_multiline():
 
 name = "Łukasz"
 (f"hello {name}", F"hello {name}")
-(b"", B"")
+(b"", b"")
 ("", "")
 (r"", R"")
 
 (rf"", fr"", Rf"", fR"", rF"", Fr"", RF"", FR"")
-(rb"", br"", Rb"", bR"", rB"", Br"", RB"", BR"")
+(rb"", rb"", Rb"", Rb"", rb"", rb"", Rb"", Rb"")
 
 
 def docstring_singleline():