Skip to content

Commit

Permalink
Preserve quotes in generated byte strings (astral-sh#15778)
Browse files Browse the repository at this point in the history
## Summary

This is a very closely related follow-up to astral-sh#15726, adding the same
quote-preserving behavior to bytestrings. Only one rule (UP018) was
affected this time, and it was easy to mirror the plain string changes.

## Test Plan

Existing tests
  • Loading branch information
ntBre authored Jan 28, 2025
1 parent 9c93844 commit 98d20a8
Show file tree
Hide file tree
Showing 5 changed files with 76 additions and 35 deletions.
6 changes: 6 additions & 0 deletions crates/ruff_linter/src/checkers/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down
15 changes: 10 additions & 5 deletions crates/ruff_linter/src/rules/pyupgrade/rules/native_literals.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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,
Expand Down
51 changes: 39 additions & 12 deletions crates/ruff_python_ast/src/nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]>,
Expand All @@ -1920,7 +1947,7 @@ impl BytesLiteral {
Self {
range,
value: Box::new([]),
flags: BytesLiteralFlags::default().with_invalid(),
flags: BytesLiteralFlags::empty().with_invalid(),
}
}
}
Expand Down Expand Up @@ -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() {
Expand Down
37 changes: 20 additions & 17 deletions crates/ruff_python_codegen/src/generator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down Expand Up @@ -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, .. }) => {
Expand Down Expand Up @@ -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}""#);
Expand Down Expand Up @@ -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]
Expand Down
2 changes: 1 addition & 1 deletion crates/ruff_python_formatter/tests/normalizer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
});
}
}
Expand Down

0 comments on commit 98d20a8

Please sign in to comment.