Add support for reformatting byte strings (#3176)

crates/ruff/src/rules/flake8_bugbear/rules/raise_without_from_inside_except.rs

@@ -1,5 +1,5 @@
 use ruff_macros::{define_violation, derive_message_formats};
-use ruff_python::string::is_lower;
+use ruff_python::str::is_lower;
 use rustpython_parser::ast::{ExprKind, Stmt, StmtKind};
 
 use crate::ast::types::Range;

crates/ruff/src/rules/flake8_simplify/rules/yoda_conditions.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use libcst_native::{Codegen, CodegenState, CompOp};
 use ruff_macros::{define_violation, derive_message_formats};
-use ruff_python::string::{self};
+use ruff_python::str::{self};
 use rustpython_parser::ast::{Cmpop, Expr, ExprKind};
 
 use crate::ast::types::Range;
@@ -46,10 +46,10 @@ impl Violation for YodaConditions {
 /// Return `true` if an [`Expr`] is a constant or a constant-like name.
 fn is_constant_like(expr: &Expr) -> bool {
     match &expr.node {
-        ExprKind::Attribute { attr, .. } => string::is_upper(attr),
+        ExprKind::Attribute { attr, .. } => str::is_upper(attr),
         ExprKind::Constant { .. } => true,
         ExprKind::Tuple { elts, .. } => elts.iter().all(is_constant_like),
-        ExprKind::Name { id, .. } => string::is_upper(id),
+        ExprKind::Name { id, .. } => str::is_upper(id),
         _ => false,
     }
 }

crates/ruff/src/rules/isort/sorting.rs

@@ -3,7 +3,7 @@ use std::cmp::Ordering;
 use std::collections::BTreeSet;
 
 use crate::rules::isort::types::Importable;
-use ruff_python::string;
+use ruff_python::str;
 
 use super::settings::RelativeImportsOrder;
 use super::types::EitherImport::{Import, ImportFrom};
@@ -31,7 +31,7 @@ fn prefix(
     } else if variables.contains(name) {
         // Ex) `variable`
         Prefix::Variables
-    } else if name.len() > 1 && string::is_upper(name) {
+    } else if name.len() > 1 && str::is_upper(name) {
         // Ex) `CONSTANT`
         Prefix::Constants
     } else if name.chars().next().map_or(false, char::is_uppercase) {

crates/ruff/src/rules/pep8_naming/helpers.rs

@@ -1,5 +1,5 @@
 use itertools::Itertools;
-use ruff_python::string::{is_lower, is_upper};
+use ruff_python::str::{is_lower, is_upper};
 use rustpython_parser::ast::{ExprKind, Stmt, StmtKind};
 
 use crate::checkers::ast::Checker;

crates/ruff/src/rules/pep8_naming/rules/camelcase_imported_as_acronym.rs

@@ -1,5 +1,5 @@
 use ruff_macros::{define_violation, derive_message_formats};
-use ruff_python::string::{self};
+use ruff_python::str::{self};
 use rustpython_parser::ast::Stmt;
 
 use crate::ast::helpers::identifier_range;
@@ -56,8 +56,8 @@ pub fn camelcase_imported_as_acronym(
     locator: &Locator,
 ) -> Option<Diagnostic> {
     if helpers::is_camelcase(name)
-        && !string::is_lower(asname)
-        && string::is_upper(asname)
+        && !str::is_lower(asname)
+        && str::is_upper(asname)
         && helpers::is_acronym(name, asname)
     {
         return Some(Diagnostic::new(

crates/ruff/src/rules/pep8_naming/rules/camelcase_imported_as_constant.rs

@@ -1,5 +1,5 @@
 use ruff_macros::{define_violation, derive_message_formats};
-use ruff_python::string::{self};
+use ruff_python::str::{self};
 use rustpython_parser::ast::Stmt;
 
 use crate::ast::helpers::identifier_range;
@@ -53,8 +53,8 @@ pub fn camelcase_imported_as_constant(
     locator: &Locator,
 ) -> Option<Diagnostic> {
     if helpers::is_camelcase(name)
-        && !string::is_lower(asname)
-        && string::is_upper(asname)
+        && !str::is_lower(asname)
+        && str::is_upper(asname)
         && !helpers::is_acronym(name, asname)
     {
         return Some(Diagnostic::new(

crates/ruff/src/rules/pep8_naming/rules/camelcase_imported_as_lowercase.rs

@@ -1,5 +1,5 @@
 use ruff_macros::{define_violation, derive_message_formats};
-use ruff_python::string;
+use ruff_python::str;
 use rustpython_parser::ast::Stmt;
 
 use crate::ast::helpers::identifier_range;
@@ -52,7 +52,7 @@ pub fn camelcase_imported_as_lowercase(
     asname: &str,
     locator: &Locator,
 ) -> Option<Diagnostic> {
-    if helpers::is_camelcase(name) && string::is_lower(asname) {
+    if helpers::is_camelcase(name) && str::is_lower(asname) {
         return Some(Diagnostic::new(
             CamelcaseImportedAsLowercase {
                 name: name.to_string(),

crates/ruff/src/rules/pep8_naming/rules/constant_imported_as_non_constant.rs

@@ -1,5 +1,5 @@
 use ruff_macros::{define_violation, derive_message_formats};
-use ruff_python::string;
+use ruff_python::str;
 use rustpython_parser::ast::Stmt;
 
 use crate::ast::helpers::identifier_range;
@@ -52,7 +52,7 @@ pub fn constant_imported_as_non_constant(
     asname: &str,
     locator: &Locator,
 ) -> Option<Diagnostic> {
-    if string::is_upper(name) && !string::is_upper(asname) {
+    if str::is_upper(name) && !str::is_upper(asname) {
         return Some(Diagnostic::new(
             ConstantImportedAsNonConstant {
                 name: name.to_string(),

crates/ruff/src/rules/pep8_naming/rules/lowercase_imported_as_non_lowercase.rs

@@ -1,5 +1,5 @@
 use ruff_macros::{define_violation, derive_message_formats};
-use ruff_python::string;
+use ruff_python::str;
 use rustpython_parser::ast::Stmt;
 
 use crate::ast::helpers::identifier_range;
@@ -51,7 +51,7 @@ pub fn lowercase_imported_as_non_lowercase(
     asname: &str,
     locator: &Locator,
 ) -> Option<Diagnostic> {
-    if !string::is_upper(name) && string::is_lower(name) && asname.to_lowercase() != asname {
+    if !str::is_upper(name) && str::is_lower(name) && asname.to_lowercase() != asname {
         return Some(Diagnostic::new(
             LowercaseImportedAsNonLowercase {
                 name: name.to_string(),

crates/ruff/src/rules/pydocstyle/helpers.rs

@@ -1,4 +1,4 @@
-use ruff_python::string::{
+use ruff_python::str::{
     SINGLE_QUOTE_PREFIXES, SINGLE_QUOTE_SUFFIXES, TRIPLE_QUOTE_PREFIXES, TRIPLE_QUOTE_SUFFIXES,
 };
 

crates/ruff/src/rules/pydocstyle/rules/multi_line_summary_start.rs

@@ -1,5 +1,5 @@
 use ruff_macros::{define_violation, derive_message_formats};
-use ruff_python::string::TRIPLE_QUOTE_PREFIXES;
+use ruff_python::str::TRIPLE_QUOTE_PREFIXES;
 
 use crate::ast::types::Range;
 use crate::ast::whitespace::LinesWithTrailingNewline;

crates/ruff/src/rules/pyflakes/fixes.rs

@@ -1,6 +1,6 @@
 use anyhow::{bail, Result};
 use libcst_native::{Call, Codegen, CodegenState, Dict, DictElement, Expression};
-use ruff_python::string::strip_quotes_and_prefixes;
+use ruff_python::str::strip_quotes_and_prefixes;
 use rustpython_parser::ast::{Excepthandler, Expr};
 use rustpython_parser::{lexer, Mode, Tok};
 

crates/ruff_python/src/bytes.rs

@@ -0,0 +1,10 @@
+/// See: <https://docs.python.org/3/reference/lexical_analysis.html#string-and-bytes-literals>
+pub const TRIPLE_QUOTE_PREFIXES: &[&str] = &[
+    "br'''", "rb'''", "bR'''", "Rb'''", "Br'''", "rB'''", "RB'''", "BR'''", "b'''", "br\"\"\"",
+    "rb\"\"\"", "bR\"\"\"", "Rb\"\"\"", "Br\"\"\"", "rB\"\"\"", "RB\"\"\"", "BR\"\"\"", "b\"\"\"",
+    "B\"\"\"",
+];
+pub const SINGLE_QUOTE_PREFIXES: &[&str] = &[
+    "br'", "rb'", "bR'", "Rb'", "Br'", "rB'", "RB'", "BR'", "b'", "br\"", "rb\"", "bR\"", "Rb\"",
+    "Br\"", "rB\"", "RB\"", "BR\"", "b\"", "B\"",
+];

crates/ruff_python/src/lib.rs

@@ -1,7 +1,8 @@
 pub mod builtins;
+pub mod bytes;
 pub mod future;
 pub mod identifiers;
 pub mod keyword;
-pub mod string;
+pub mod str;
 pub mod sys;
 pub mod typing;

crates/ruff_python/src/str.rs

@@ -6,7 +6,7 @@ pub const TRIPLE_QUOTE_PREFIXES: &[&str] = &[
     "u\"\"\"", "u'''", "r\"\"\"", "r'''", "U\"\"\"", "U'''", "R\"\"\"", "R'''", "\"\"\"", "'''",
 ];
 pub const SINGLE_QUOTE_PREFIXES: &[&str] = &[
-    "u\"", "u'", "r\"", "r'", "u\"", "u'", "r\"", "r'", "U\"", "U'", "R\"", "R'", "\"", "'",
+    "u\"", "u'", "r\"", "r'", "U\"", "U'", "R\"", "R'", "\"", "'",
 ];
 pub const TRIPLE_QUOTE_SUFFIXES: &[&str] = &["\"\"\"", "'''"];
 pub const SINGLE_QUOTE_SUFFIXES: &[&str] = &["\"", "'"];
@@ -53,7 +53,7 @@ pub fn strip_quotes_and_prefixes(s: &str) -> &str {
 
 #[cfg(test)]
 mod tests {
-    use crate::string::{is_lower, is_upper, strip_quotes_and_prefixes};
+    use crate::str::{is_lower, is_upper, strip_quotes_and_prefixes};
 
     #[test]
     fn test_is_lower() {

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():