litellm/litellm-rust/crates/python-compat/src/repr.rs

//! `repr(value)` and `str(value)` byte for byte.
//!
//! LiteLLM hashes `str(value)` into cache keys and writes `str(dict)` into Redis, so these
//! strings are persisted identifiers rather than display text: every quote choice, escape,
//! and float digit must match CPython.

use std::fmt::Write;

use crate::Value;

/// Generated by `scripts/generate_nonprintable.py`; see `generated/nonprintable.rs`.
mod nonprintable {
    include!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/generated/nonprintable.rs"
    ));
}

/// Unicode version of the printable-character table, from the Python that generated it.
pub const UNICODE_VERSION: &str = nonprintable::UNICODE_VERSION;

/// `repr(value)`.
pub fn repr(value: &Value) -> String {
    let mut out = String::new();
    write_repr(&mut out, value);
    out
}

/// `str(value)`: a string's own contents, and `repr` for every other value.
pub fn to_str(value: &Value) -> String {
    match value {
        Value::Str(text) => text.clone(),
        value => repr(value),
    }
}

/// `repr(float)`, the shortest round-trip form with CPython's exponent thresholds.
pub fn float_repr(value: f64) -> String {
    format_float(value, true)
}

fn write_repr(out: &mut String, value: &Value) {
    match value {
        Value::None => out.push_str("None"),
        Value::Bool(true) => out.push_str("True"),
        Value::Bool(false) => out.push_str("False"),
        Value::Int(value) => {
            let _ = write!(out, "{value}");
        }
        Value::Float(value) => out.push_str(&float_repr(*value)),
        Value::Complex { re, im } => write_complex(out, *re, *im),
        Value::Str(text) => write_str(out, text),
        Value::Bytes(bytes) => write_bytes(out, bytes),
        Value::List(values) => write_sequence(out, '[', values, ']'),
        Value::Tuple(values) if values.len() == 1 => {
            out.push('(');
            write_repr(out, &values[0]);
            out.push_str(",)");
        }
        Value::Tuple(values) => write_sequence(out, '(', values, ')'),
        Value::Set(values) if values.is_empty() => out.push_str("set()"),
        Value::Set(values) => write_sequence(out, '{', values, '}'),
        Value::Dict(entries) => {
            out.push('{');
            for (index, (key, value)) in entries.iter().enumerate() {
                if index > 0 {
                    out.push_str(", ");
                }
                write_repr(out, key);
                out.push_str(": ");
                write_repr(out, value);
            }
            out.push('}');
        }
    }
}

fn write_sequence(out: &mut String, open: char, values: &[Value], close: char) {
    out.push(open);
    for (index, value) in values.iter().enumerate() {
        if index > 0 {
            out.push_str(", ");
        }
        write_repr(out, value);
    }
    out.push(close);
}

/// `complex.__repr__`: a `+0.0` real part prints only the imaginary part, without parens.
fn write_complex(out: &mut String, re: f64, im: f64) {
    if re == 0.0 && re.is_sign_positive() {
        out.push_str(&format_float(im, false));
        out.push('j');
        return;
    }
    out.push('(');
    out.push_str(&format_float(re, false));
    let im_text = format_float(im, false);
    if !im_text.starts_with('-') {
        out.push('+');
    }
    out.push_str(&im_text);
    out.push_str("j)");
}

/// `PyOS_double_to_string(value, 'r', 0, flags)`: scientific notation below 1e-4 and from
/// 1e16 up, with a signed exponent of at least two digits. `add_dot_0` is
/// `Py_DTSF_ADD_DOT_0`, which `float` sets and `complex` does not.
fn format_float(value: f64, add_dot_0: bool) -> String {
    if value.is_nan() {
        return "nan".to_owned();
    }
    if value.is_infinite() {
        return if value > 0.0 { "inf" } else { "-inf" }.to_owned();
    }
    // Rust's `{:e}` prints the shortest round-trip digits, like CPython's 'r' mode.
    let scientific = format!("{value:e}");
    let (mantissa, exponent) = scientific
        .split_once('e')
        .expect("`{:e}` always prints an exponent");
    let exponent: i32 = exponent.parse().expect("`{:e}` exponent is an integer");
    let (sign, mantissa) = match mantissa.strip_prefix('-') {
        Some(mantissa) => ("-", mantissa),
        None => ("", mantissa),
    };
    let digits: String = mantissa.chars().filter(|ch| *ch != '.').collect();

    let mut out = String::from(sign);
    if !(-4..16).contains(&exponent) {
        out.push_str(&digits[..1]);
        if digits.len() > 1 {
            out.push('.');
            out.push_str(&digits[1..]);
        }
        let _ = write!(
            out,
            "e{}{:02}",
            if exponent < 0 { '-' } else { '+' },
            exponent.unsigned_abs()
        );
    } else if exponent < 0 {
        out.push_str("0.");
        out.extend(std::iter::repeat_n('0', (-exponent - 1) as usize));
        out.push_str(&digits);
    } else {
        let integer_digits = exponent as usize + 1;
        if digits.len() > integer_digits {
            out.push_str(&digits[..integer_digits]);
            out.push('.');
            out.push_str(&digits[integer_digits..]);
        } else {
            out.push_str(&digits);
            out.extend(std::iter::repeat_n('0', integer_digits - digits.len()));
            if add_dot_0 {
                out.push_str(".0");
            }
        }
    }
    out
}

/// `unicode_repr`: single quotes unless the text has a `'` and no `"`. Printable non-ASCII
/// stays literal; everything `str.isprintable()` rejects is escaped.
fn write_str(out: &mut String, text: &str) {
    let quote = if text.contains('\'') && !text.contains('"') {
        '"'
    } else {
        '\''
    };
    out.push(quote);
    for ch in text.chars() {
        match ch {
            '\\' => out.push_str("\\\\"),
            '\t' => out.push_str("\\t"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            ch if ch == quote => {
                out.push('\\');
                out.push(ch);
            }
            ch if is_printable(ch) => out.push(ch),
            ch => {
                let code = ch as u32;
                let _ = match code {
                    0..=0xff => write!(out, "\\x{code:02x}"),
                    0x100..=0xffff => write!(out, "\\u{code:04x}"),
                    _ => write!(out, "\\U{code:08x}"),
                };
            }
        }
    }
    out.push(quote);
}

/// `bytes.__repr__`: the same quote rule as `str`, with every byte outside printable ASCII
/// escaped as `\xhh`.
fn write_bytes(out: &mut String, bytes: &[u8]) {
    let quote = if bytes.contains(&b'\'') && !bytes.contains(&b'"') {
        b'"'
    } else {
        b'\''
    };
    out.push('b');
    out.push(quote as char);
    for &byte in bytes {
        match byte {
            b'\\' => out.push_str("\\\\"),
            b'\t' => out.push_str("\\t"),
            b'\n' => out.push_str("\\n"),
            b'\r' => out.push_str("\\r"),
            byte if byte == quote => {
                out.push('\\');
                out.push(byte as char);
            }
            0x20..=0x7e => out.push(byte as char),
            byte => {
                let _ = write!(out, "\\x{byte:02x}");
            }
        }
    }
    out.push(quote as char);
}

fn is_printable(ch: char) -> bool {
    let code = ch as u32;
    nonprintable::NONPRINTABLE
        .binary_search_by(|&(low, high)| {
            if high < code {
                std::cmp::Ordering::Less
            } else if low > code {
                std::cmp::Ordering::Greater
            } else {
                std::cmp::Ordering::Equal
            }
        })
        .is_err()
}