Merge pull request #4951 from youknowone/repr-quote

Port configurable repr quote from Ruff and refactoring

common/src/bytes.rs

@@ -1,62 +0,0 @@
-use crate::str::ReprOverflowError;
-
-pub fn repr(b: &[u8]) -> Result<String, ReprOverflowError> {
-    repr_with(b, &[], "")
-}
-
-pub fn repr_with(b: &[u8], prefixes: &[&str], suffix: &str) -> Result<String, ReprOverflowError> {
-    use std::fmt::Write;
-
-    let mut out_len = 0usize;
-    let mut squote = 0;
-    let mut dquote = 0;
-
-    for &ch in b {
-        let incr = match ch {
-            b'\'' => {
-                squote += 1;
-                1
-            }
-            b'"' => {
-                dquote += 1;
-                1
-            }
-            b'\\' | b'\t' | b'\r' | b'\n' => 2,
-            0x20..=0x7e => 1,
-            _ => 4, // \xHH
-        };
-        out_len = out_len.checked_add(incr).ok_or(ReprOverflowError)?;
-    }
-
-    let (quote, num_escaped_quotes) = crate::str::choose_quotes_for_repr(squote, dquote);
-    // we'll be adding backslashes in front of the existing inner quotes
-    out_len += num_escaped_quotes;
-
-    // 3 is for b prefix + outer quotes
-    out_len += 3 + prefixes.iter().map(|s| s.len()).sum::<usize>() + suffix.len();
-
-    let mut res = String::with_capacity(out_len);
-    res.extend(prefixes.iter().copied());
-    res.push('b');
-    res.push(quote);
-    for &ch in b {
-        match ch {
-            b'\t' => res.push_str("\\t"),
-            b'\n' => res.push_str("\\n"),
-            b'\r' => res.push_str("\\r"),
-            // printable ascii range
-            0x20..=0x7e => {
-                let ch = ch as char;
-                if ch == quote || ch == '\\' {
-                    res.push('\\');
-                }
-                res.push(ch);
-            }
-            _ => write!(res, "\\x{ch:02x}").unwrap(),
-        }
-    }
-    res.push(quote);
-    res.push_str(suffix);
-
-    Ok(res)
-}

common/src/escape.rs

@@ -0,0 +1,414 @@
+#[derive(Debug, Clone, Copy)]
+pub enum Quote {
+    Single,
+    Double,
+}
+
+impl Quote {
+    #[inline]
+    pub const fn swap(self) -> Quote {
+        match self {
+            Quote::Single => Quote::Double,
+            Quote::Double => Quote::Single,
+        }
+    }
+
+    #[inline]
+    pub const fn to_byte(&self) -> u8 {
+        match self {
+            Quote::Single => b'\'',
+            Quote::Double => b'"',
+        }
+    }
+
+    #[inline]
+    pub const fn to_char(&self) -> char {
+        match self {
+            Quote::Single => '\'',
+            Quote::Double => '"',
+        }
+    }
+}
+
+pub struct EscapeLayout {
+    pub quote: Quote,
+    pub len: Option<usize>,
+}
+
+pub trait Escape {
+    fn source_len(&self) -> usize;
+    fn layout(&self) -> &EscapeLayout;
+    fn changed(&self) -> bool {
+        self.layout().len != Some(self.source_len())
+    }
+
+    fn write_source(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result;
+    fn write_body_slow(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result;
+    fn write_body(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result {
+        if self.changed() {
+            self.write_body_slow(formatter)
+        } else {
+            self.write_source(formatter)
+        }
+    }
+}
+
+/// Returns the outer quotes to use and the number of quotes that need to be
+/// escaped.
+pub(crate) const fn choose_quote(
+    single_count: usize,
+    double_count: usize,
+    preferred_quote: Quote,
+) -> (Quote, usize) {
+    let (primary_count, secondary_count) = match preferred_quote {
+        Quote::Single => (single_count, double_count),
+        Quote::Double => (double_count, single_count),
+    };
+
+    // always use primary unless we have primary but no seconday
+    let use_secondary = primary_count > 0 && secondary_count == 0;
+    if use_secondary {
+        (preferred_quote.swap(), secondary_count)
+    } else {
+        (preferred_quote, primary_count)
+    }
+}
+
+pub struct UnicodeEscape<'a> {
+    source: &'a str,
+    layout: EscapeLayout,
+}
+
+impl<'a> UnicodeEscape<'a> {
+    pub fn with_forced_quote(source: &'a str, quote: Quote) -> Self {
+        let layout = EscapeLayout { quote, len: None };
+        Self { source, layout }
+    }
+    pub fn new_repr(source: &'a str) -> Self {
+        let layout = Self::repr_layout(source, Quote::Single);
+        Self { source, layout }
+    }
+
+    pub fn str_repr<'r>(&'a self) -> StrRepr<'r, 'a> {
+        StrRepr(self)
+    }
+}
+
+pub struct StrRepr<'r, 'a>(&'r UnicodeEscape<'a>);
+
+impl StrRepr<'_, '_> {
+    pub fn write(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result {
+        let quote = self.0.layout().quote.to_char();
+        formatter.write_char(quote)?;
+        self.0.write_body(formatter)?;
+        formatter.write_char(quote)
+    }
+
+    pub fn to_string(&self) -> Option<String> {
+        let mut s = String::with_capacity(self.0.layout().len?);
+        self.write(&mut s).unwrap();
+        Some(s)
+    }
+}
+
+impl std::fmt::Display for StrRepr<'_, '_> {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.write(formatter)
+    }
+}
+
+impl UnicodeEscape<'_> {
+    const REPR_RESERVED_LEN: usize = 2; // for quotes
+
+    pub fn repr_layout(source: &str, preferred_quote: Quote) -> EscapeLayout {
+        Self::output_layout_with_checker(source, preferred_quote, |a, b| {
+            Some((a as isize).checked_add(b as isize)? as usize)
+        })
+    }
+
+    fn output_layout_with_checker(
+        source: &str,
+        preferred_quote: Quote,
+        length_add: impl Fn(usize, usize) -> Option<usize>,
+    ) -> EscapeLayout {
+        let mut out_len = Self::REPR_RESERVED_LEN;
+        let mut single_count = 0;
+        let mut double_count = 0;
+
+        for ch in source.chars() {
+            let incr = match ch {
+                '\'' => {
+                    single_count += 1;
+                    1
+                }
+                '"' => {
+                    double_count += 1;
+                    1
+                }
+                c => Self::escaped_char_len(c),
+            };
+            let Some(new_len) = length_add(out_len, incr) else {
+                #[cold]
+                fn stop(single_count: usize, double_count: usize, preferred_quote: Quote) -> EscapeLayout {
+                    EscapeLayout { quote: choose_quote(single_count, double_count, preferred_quote).0, len: None }
+                }
+                return stop(single_count, double_count, preferred_quote);
+            };
+            out_len = new_len;
+        }
+
+        let (quote, num_escaped_quotes) = choose_quote(single_count, double_count, preferred_quote);
+        // we'll be adding backslashes in front of the existing inner quotes
+        let Some(out_len) = length_add(out_len, num_escaped_quotes) else {
+            return EscapeLayout { quote, len: None };
+        };
+
+        EscapeLayout {
+            quote,
+            len: Some(out_len - Self::REPR_RESERVED_LEN),
+        }
+    }
+
+    fn escaped_char_len(ch: char) -> usize {
+        match ch {
+            '\\' | '\t' | '\r' | '\n' => 2,
+            ch if ch < ' ' || ch as u32 == 0x7f => 4, // \xHH
+            ch if ch.is_ascii() => 1,
+            ch if crate::char::is_printable(ch) => {
+                // max = std::cmp::max(ch, max);
+                ch.len_utf8()
+            }
+            ch if (ch as u32) < 0x100 => 4,   // \xHH
+            ch if (ch as u32) < 0x10000 => 6, // \uHHHH
+            _ => 10,                          // \uHHHHHHHH
+        }
+    }
+
+    fn write_char(
+        ch: char,
+        quote: Quote,
+        formatter: &mut impl std::fmt::Write,
+    ) -> std::fmt::Result {
+        match ch {
+            '\n' => formatter.write_str("\\n"),
+            '\t' => formatter.write_str("\\t"),
+            '\r' => formatter.write_str("\\r"),
+            // these 2 branches *would* be handled below, but we shouldn't have to do a
+            // unicodedata lookup just for ascii characters
+            '\x20'..='\x7e' => {
+                // printable ascii range
+                if ch == quote.to_char() || ch == '\\' {
+                    formatter.write_char('\\')?;
+                }
+                formatter.write_char(ch)
+            }
+            ch if ch.is_ascii() => {
+                write!(formatter, "\\x{:02x}", ch as u8)
+            }
+            ch if crate::char::is_printable(ch) => formatter.write_char(ch),
+            '\0'..='\u{ff}' => {
+                write!(formatter, "\\x{:02x}", ch as u32)
+            }
+            '\0'..='\u{ffff}' => {
+                write!(formatter, "\\u{:04x}", ch as u32)
+            }
+            _ => {
+                write!(formatter, "\\U{:08x}", ch as u32)
+            }
+        }
+    }
+}
+
+impl<'a> Escape for UnicodeEscape<'a> {
+    fn source_len(&self) -> usize {
+        self.source.len()
+    }
+
+    fn layout(&self) -> &EscapeLayout {
+        &self.layout
+    }
+
+    fn write_source(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result {
+        formatter.write_str(self.source)
+    }
+
+    #[cold]
+    fn write_body_slow(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result {
+        for ch in self.source.chars() {
+            Self::write_char(ch, self.layout().quote, formatter)?;
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod unicode_escapse_tests {
+    use super::*;
+
+    #[test]
+    fn changed() {
+        fn test(s: &str) -> bool {
+            UnicodeEscape::new_repr(s).changed()
+        }
+        assert!(!test("hello"));
+        assert!(!test("'hello'"));
+        assert!(!test("\"hello\""));
+
+        assert!(test("'\"hello"));
+        assert!(test("hello\n"));
+    }
+}
+
+pub struct AsciiEscape<'a> {
+    source: &'a [u8],
+    layout: EscapeLayout,
+}
+
+impl<'a> AsciiEscape<'a> {
+    pub fn new(source: &'a [u8], layout: EscapeLayout) -> Self {
+        Self { source, layout }
+    }
+    pub fn with_forced_quote(source: &'a [u8], quote: Quote) -> Self {
+        let layout = EscapeLayout { quote, len: None };
+        Self { source, layout }
+    }
+    pub fn new_repr(source: &'a [u8]) -> Self {
+        let layout = Self::repr_layout(source, Quote::Single);
+        Self { source, layout }
+    }
+
+    pub fn bytes_repr<'r>(&'a self) -> BytesRepr<'r, 'a> {
+        BytesRepr(self)
+    }
+}
+
+impl AsciiEscape<'_> {
+    pub fn repr_layout(source: &[u8], preferred_quote: Quote) -> EscapeLayout {
+        Self::output_layout_with_checker(source, preferred_quote, 3, |a, b| {
+            Some((a as isize).checked_add(b as isize)? as usize)
+        })
+    }
+
+    pub fn named_repr_layout(source: &[u8], name: &str) -> EscapeLayout {
+        Self::output_layout_with_checker(source, Quote::Single, name.len() + 2 + 3, |a, b| {
+            Some((a as isize).checked_add(b as isize)? as usize)
+        })
+    }
+
+    fn output_layout_with_checker(
+        source: &[u8],
+        preferred_quote: Quote,
+        reserved_len: usize,
+        length_add: impl Fn(usize, usize) -> Option<usize>,
+    ) -> EscapeLayout {
+        let mut out_len = reserved_len;
+        let mut single_count = 0;
+        let mut double_count = 0;
+
+        for ch in source.iter() {
+            let incr = match ch {
+                b'\'' => {
+                    single_count += 1;
+                    1
+                }
+                b'"' => {
+                    double_count += 1;
+                    1
+                }
+                c => Self::escaped_char_len(*c),
+            };
+            let Some(new_len) = length_add(out_len, incr) else {
+                #[cold]
+                fn stop(single_count: usize, double_count: usize, preferred_quote: Quote) -> EscapeLayout {
+                    EscapeLayout { quote: choose_quote(single_count, double_count, preferred_quote).0, len: None }
+                }
+                return stop(single_count, double_count, preferred_quote);
+            };
+            out_len = new_len;
+        }
+
+        let (quote, num_escaped_quotes) = choose_quote(single_count, double_count, preferred_quote);
+        // we'll be adding backslashes in front of the existing inner quotes
+        let Some(out_len) = length_add(out_len, num_escaped_quotes) else {
+            return EscapeLayout { quote, len: None };
+        };
+
+        EscapeLayout {
+            quote,
+            len: Some(out_len - reserved_len),
+        }
+    }
+
+    fn escaped_char_len(ch: u8) -> usize {
+        match ch {
+            b'\\' | b'\t' | b'\r' | b'\n' => 2,
+            0x20..=0x7e => 1,
+            _ => 4, // \xHH
+        }
+    }
+
+    fn write_char(ch: u8, quote: Quote, formatter: &mut impl std::fmt::Write) -> std::fmt::Result {
+        match ch {
+            b'\t' => formatter.write_str("\\t"),
+            b'\n' => formatter.write_str("\\n"),
+            b'\r' => formatter.write_str("\\r"),
+            0x20..=0x7e => {
+                // printable ascii range
+                if ch == quote.to_byte() || ch == b'\\' {
+                    formatter.write_char('\\')?;
+                }
+                formatter.write_char(ch as char)
+            }
+            ch => write!(formatter, "\\x{ch:02x}"),
+        }
+    }
+}
+
+impl<'a> Escape for AsciiEscape<'a> {
+    fn source_len(&self) -> usize {
+        self.source.len()
+    }
+
+    fn layout(&self) -> &EscapeLayout {
+        &self.layout
+    }
+
+    fn write_source(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result {
+        formatter.write_str(unsafe {
+            // SAFETY: this function must be called only when source is printable ascii characters
+            std::str::from_utf8_unchecked(self.source)
+        })
+    }
+
+    #[cold]
+    fn write_body_slow(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result {
+        for ch in self.source.iter() {
+            Self::write_char(*ch, self.layout().quote, formatter)?;
+        }
+        Ok(())
+    }
+}
+
+pub struct BytesRepr<'r, 'a>(&'r AsciiEscape<'a>);
+
+impl BytesRepr<'_, '_> {
+    pub fn write(&self, formatter: &mut impl std::fmt::Write) -> std::fmt::Result {
+        let quote = self.0.layout().quote.to_char();
+        formatter.write_char('b')?;
+        formatter.write_char(quote)?;
+        self.0.write_body(formatter)?;
+        formatter.write_char(quote)
+    }
+
+    pub fn to_string(&self) -> Option<String> {
+        let mut s = String::with_capacity(self.0.layout().len?);
+        self.write(&mut s).unwrap();
+        Some(s)
+    }
+}
+
+impl std::fmt::Display for BytesRepr<'_, '_> {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.write(formatter)
+    }
+}

common/src/lib.rs

@@ -7,13 +7,13 @@ pub use macros::*;
 pub mod atomic;
 pub mod borrow;
 pub mod boxvec;
-pub mod bytes;
 pub mod cformat;
 pub mod char;
 pub mod cmp;
 #[cfg(any(unix, windows, target_os = "wasi"))]
 pub mod crt_fd;
 pub mod encodings;
+pub mod escape;
 pub mod float_ops;
 pub mod format;
 pub mod hash;

common/src/str.rs

@@ -3,11 +3,7 @@ use crate::{
     hash::PyHash,
 };
 use ascii::AsciiString;
-use once_cell::unsync::OnceCell;
-use std::{
-    fmt,
-    ops::{Bound, RangeBounds},
-};
+use std::ops::{Bound, RangeBounds};
 
 #[cfg(not(target_arch = "wasm32"))]
 #[allow(non_camel_case_types)]
@@ -342,158 +338,6 @@ macro_rules! ascii {
     }};
 }
 
-/// Get a Display-able type that formats to the python `repr()` of the string value
-#[inline]
-pub fn repr(s: &str) -> Repr<'_> {
-    Repr {
-        s,
-        info: OnceCell::new(),
-    }
-}
-
-#[derive(Debug, Copy, Clone)]
-#[non_exhaustive]
-pub struct ReprOverflowError;
-impl fmt::Display for ReprOverflowError {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        f.write_str("string is too long to generate repr")
-    }
-}
-
-#[derive(Copy, Clone)]
-struct ReprInfo {
-    dquoted: bool,
-    out_len: usize,
-}
-impl ReprInfo {
-    fn get(s: &str) -> Result<Self, ReprOverflowError> {
-        let mut out_len = 0usize;
-        let mut squote = 0;
-        let mut dquote = 0;
-
-        for ch in s.chars() {
-            let incr = match ch {
-                '\'' => {
-                    squote += 1;
-                    1
-                }
-                '"' => {
-                    dquote += 1;
-                    1
-                }
-                '\\' | '\t' | '\r' | '\n' => 2,
-                ch if ch < ' ' || ch as u32 == 0x7f => 4, // \xHH
-                ch if ch.is_ascii() => 1,
-                ch if crate::char::is_printable(ch) => {
-                    // max = std::cmp::max(ch, max);
-                    ch.len_utf8()
-                }
-                ch if (ch as u32) < 0x100 => 4,   // \xHH
-                ch if (ch as u32) < 0x10000 => 6, // \uHHHH
-                _ => 10,                          // \uHHHHHHHH
-            };
-            out_len += incr;
-            if out_len > std::isize::MAX as usize {
-                return Err(ReprOverflowError);
-            }
-        }
-
-        let (quote, num_escaped_quotes) = choose_quotes_for_repr(squote, dquote);
-        // we'll be adding backslashes in front of the existing inner quotes
-        out_len += num_escaped_quotes;
-
-        // start and ending quotes
-        out_len += 2;
-
-        let dquoted = quote == '"';
-
-        Ok(ReprInfo { dquoted, out_len })
-    }
-}
-
-pub struct Repr<'a> {
-    s: &'a str,
-    // the tuple is dquouted, out_len
-    info: OnceCell<Result<ReprInfo, ReprOverflowError>>,
-}
-impl Repr<'_> {
-    fn get_info(&self) -> Result<ReprInfo, ReprOverflowError> {
-        *self.info.get_or_init(|| ReprInfo::get(self.s))
-    }
-
-    /// Same as `<Self as ToString>::to_string()`, but checks for a possible OverflowError.
-    pub fn to_string_checked(&self) -> Result<String, ReprOverflowError> {
-        let info = self.get_info()?;
-        let mut repr = String::with_capacity(info.out_len);
-        self._fmt(&mut repr, info).unwrap();
-        Ok(repr)
-    }
-
-    fn _fmt<W: fmt::Write>(&self, repr: &mut W, info: ReprInfo) -> fmt::Result {
-        let s = self.s;
-        let in_len = s.len();
-        let ReprInfo { dquoted, out_len } = info;
-
-        let quote = if dquoted { '"' } else { '\'' };
-        // if we don't need to escape anything we can just copy
-        let unchanged = out_len == in_len;
-
-        repr.write_char(quote)?;
-        if unchanged {
-            repr.write_str(s)?;
-        } else {
-            for ch in s.chars() {
-                match ch {
-                    '\n' => repr.write_str("\\n"),
-                    '\t' => repr.write_str("\\t"),
-                    '\r' => repr.write_str("\\r"),
-                    // these 2 branches *would* be handled below, but we shouldn't have to do a
-                    // unicodedata lookup just for ascii characters
-                    '\x20'..='\x7e' => {
-                        // printable ascii range
-                        if ch == quote || ch == '\\' {
-                            repr.write_char('\\')?;
-                        }
-                        repr.write_char(ch)
-                    }
-                    ch if ch.is_ascii() => {
-                        write!(repr, "\\x{:02x}", ch as u8)
-                    }
-                    ch if crate::char::is_printable(ch) => repr.write_char(ch),
-                    '\0'..='\u{ff}' => {
-                        write!(repr, "\\x{:02x}", ch as u32)
-                    }
-                    '\0'..='\u{ffff}' => {
-                        write!(repr, "\\u{:04x}", ch as u32)
-                    }
-                    _ => {
-                        write!(repr, "\\U{:08x}", ch as u32)
-                    }
-                }?;
-            }
-        }
-        repr.write_char(quote)
-    }
-}
-
-impl fmt::Display for Repr<'_> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let info = self.get_info().unwrap();
-        self._fmt(f, info)
-    }
-}
-
-/// returns the outer quotes to use and the number of quotes that need to be escaped
-pub(crate) fn choose_quotes_for_repr(num_squotes: usize, num_dquotes: usize) -> (char, usize) {
-    // always use squote unless we have squotes but no dquotes
-    let use_dquote = num_squotes > 0 && num_dquotes == 0;
-    if use_dquote {
-        ('"', num_dquotes)
-    } else {
-        ('\'', num_squotes)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;