use std::collections::HashMap;
use std::fmt::Write;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::OnceLock;
const TABLES_JSON: &str = include_str!("../assets/spec/pycompat-tables.json");
pub const SURROGATE_SENTINEL_BASE: u32 = 0x10FC00;
static STDIN_SURROGATES: AtomicBool = AtomicBool::new(false);
pub fn surrogate_sentinels_active() -> bool {
STDIN_SURROGATES.load(Ordering::Relaxed)
}
pub fn set_surrogate_sentinels_active(on: bool) {
STDIN_SURROGATES.store(on, Ordering::Relaxed);
}
pub fn sentinel_surrogate(c: char) -> Option<u32> {
let cp = c as u32;
if surrogate_sentinels_active() && (SURROGATE_SENTINEL_BASE + 0x80..=SURROGATE_SENTINEL_BASE + 0xFF).contains(&cp) {
Some(0xDC00 + (cp - SURROGATE_SENTINEL_BASE))
} else {
None
}
}
pub fn decode_stdin_surrogateescape(bytes: &[u8]) -> String {
let mut out = String::with_capacity(bytes.len());
let mut rest = bytes;
loop {
match std::str::from_utf8(rest) {
Ok(s) => {
out.push_str(s);
return out;
}
Err(e) => {
let valid = e.valid_up_to();
out.push_str(std::str::from_utf8(&rest[..valid]).expect("valid prefix"));
let bad = e.error_len().unwrap_or(rest.len() - valid);
for &b in &rest[valid..valid + bad] {
out.push(
char::from_u32(SURROGATE_SENTINEL_BASE + b as u32)
.expect("plane-16 PUA sentinel"),
);
STDIN_SURROGATES.store(true, Ordering::Relaxed);
}
rest = &rest[valid + bad..];
}
}
}
}
pub fn encode_stdout_surrogateescape(text: &str) -> std::borrow::Cow<'_, [u8]> {
if !surrogate_sentinels_active() || !text.chars().any(|c| sentinel_surrogate(c).is_some()) {
return std::borrow::Cow::Borrowed(text.as_bytes());
}
let mut out = Vec::with_capacity(text.len());
for c in text.chars() {
if let Some(sur) = sentinel_surrogate(c) {
out.push((sur - 0xDC00) as u8);
} else {
let mut buf = [0u8; 4];
out.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
}
}
std::borrow::Cow::Owned(out)
}
struct Tables {
casefold: HashMap<u32, String>,
str_whitespace: Vec<(u32, u32)>,
splitlines_boundaries: Vec<u32>,
isprintable: Vec<(u32, u32)>,
re_digit: Vec<(u32, u32)>,
re_word: Vec<(u32, u32)>,
}
fn parse_ranges(v: &serde_json::Value) -> Vec<(u32, u32)> {
v.as_array()
.expect("range table must be an array")
.iter()
.map(|pair| {
let p = pair.as_array().expect("range entry must be a pair");
(
p[0].as_u64().expect("range start") as u32,
p[1].as_u64().expect("range end") as u32,
)
})
.collect()
}
fn tables() -> &'static Tables {
static TABLES: OnceLock<Tables> = OnceLock::new();
TABLES.get_or_init(|| {
let root: serde_json::Value =
serde_json::from_str(TABLES_JSON).expect("pycompat-tables.json must parse");
let casefold = root["casefold"]
.as_object()
.expect("casefold table")
.iter()
.map(|(k, v)| {
(
k.parse::<u32>().expect("casefold key"),
v.as_str().expect("casefold value").to_string(),
)
})
.collect();
let splitlines_boundaries = root["splitlines_boundaries"]
.as_array()
.expect("splitlines_boundaries")
.iter()
.map(|v| v.as_u64().expect("boundary cp") as u32)
.collect();
Tables {
casefold,
str_whitespace: parse_ranges(&root["str_whitespace"]),
splitlines_boundaries,
isprintable: parse_ranges(&root["isprintable"]),
re_digit: parse_ranges(&root["re_digit"]),
re_word: parse_ranges(&root["re_word"]),
}
})
}
fn in_ranges(ranges: &[(u32, u32)], cp: u32) -> bool {
let idx = ranges.partition_point(|&(start, _)| start <= cp);
idx > 0 && cp <= ranges[idx - 1].1
}
pub fn py_casefold(s: &str) -> String {
let t = tables();
let mut out = String::with_capacity(s.len());
for c in s.chars() {
match t.casefold.get(&(c as u32)) {
Some(folded) => out.push_str(folded),
None => out.push(c),
}
}
out
}
pub fn py_is_space(c: char) -> bool {
in_ranges(&tables().str_whitespace, c as u32)
}
pub fn py_strip(s: &str) -> &str {
py_rstrip(py_lstrip(s))
}
pub fn py_lstrip(s: &str) -> &str {
s.trim_start_matches(py_is_space)
}
pub fn py_rstrip(s: &str) -> &str {
s.trim_end_matches(py_is_space)
}
fn is_line_boundary(c: char) -> bool {
let b = &tables().splitlines_boundaries;
b.binary_search(&(c as u32)).is_ok()
}
pub fn py_splitlines(s: &str) -> Vec<&str> {
let mut out = Vec::new();
let mut start = 0usize;
let mut iter = s.char_indices().peekable();
while let Some((i, c)) = iter.next() {
if is_line_boundary(c) {
out.push(&s[start..i]);
let mut end = i + c.len_utf8();
if c == '\r' {
if let Some(&(j, '\n')) = iter.peek() {
iter.next();
end = j + 1;
}
}
start = end;
}
}
if start < s.len() {
out.push(&s[start..]);
}
out
}
pub fn first_nonempty_line(s: &str) -> &str {
py_splitlines(s)
.into_iter()
.map(py_strip)
.find(|l| !l.is_empty())
.unwrap_or("")
}
pub fn read_text_universal(path: &str) -> Option<String> {
let bytes = std::fs::read(path).ok()?;
let text = String::from_utf8(bytes).ok()?;
Some(text.replace("\r\n", "\n").replace('\r', "\n"))
}
pub fn py_is_printable(c: char) -> bool {
in_ranges(&tables().isprintable, c as u32)
}
pub fn is_re_digit(c: char) -> bool {
in_ranges(&tables().re_digit, c as u32)
}
pub fn is_re_word(c: char) -> bool {
in_ranges(&tables().re_word, c as u32)
}
pub fn py_repr_str(s: &str) -> String {
let quote = if s.contains('\'') && !s.contains('"') {
'"'
} else {
'\''
};
let mut out = String::with_capacity(s.len() + 2);
out.push(quote);
for c in s.chars() {
if c == quote || c == '\\' {
out.push('\\');
out.push(c);
} else if c == '\t' {
out.push_str("\\t");
} else if c == '\n' {
out.push_str("\\n");
} else if c == '\r' {
out.push_str("\\r");
} else if let Some(sur) = sentinel_surrogate(c) {
write!(out, "\\u{sur:04x}").unwrap();
} else if py_is_printable(c) {
out.push(c);
} else {
let cp = c as u32;
if cp < 0x100 {
write!(out, "\\x{cp:02x}").unwrap();
} else if cp < 0x10000 {
write!(out, "\\u{cp:04x}").unwrap();
} else {
write!(out, "\\U{cp:08x}").unwrap();
}
}
}
out.push(quote);
out
}
pub fn quote_plus(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for b in s.bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'_' | b'.' | b'-' | b'~' => {
out.push(b as char)
}
b' ' => out.push('+'),
_ => {
write!(out, "%{b:02X}").unwrap();
}
}
}
out
}
pub fn quote_plus_urlencode(pairs: &[(&str, &str)]) -> String {
pairs
.iter()
.map(|(k, v)| format!("{}={}", quote_plus(k), quote_plus(v)))
.collect::<Vec<_>>()
.join("&")
}
pub fn py_float_repr(x: f64) -> String {
if x.is_nan() {
return "nan".to_string();
}
if x.is_infinite() {
return if x > 0.0 { "inf" } else { "-inf" }.to_string();
}
let neg = x.is_sign_negative();
let sign = if neg { "-" } else { "" };
if x == 0.0 {
return format!("{sign}0.0");
}
let ax = x.abs();
let (n, k) = exact_decimal(ax);
let exact_decpt = n.len() as i64 - k; let (digits, decpt) = shortest_digits(ax, &n, exact_decpt);
if decpt <= -4 || decpt > 16 {
let e10 = decpt - 1;
let mantissa = if digits.len() > 1 {
format!("{}.{}", &digits[..1], &digits[1..])
} else {
digits.clone()
};
let (esign, eabs) = if e10 < 0 { ('-', -e10) } else { ('+', e10) };
format!("{sign}{mantissa}e{esign}{eabs:02}")
} else if decpt <= 0 {
format!("{sign}0.{}{}", "0".repeat((-decpt) as usize), digits)
} else if (decpt as usize) >= digits.len() {
format!(
"{sign}{}{}.0",
digits,
"0".repeat(decpt as usize - digits.len())
)
} else {
let d = decpt as usize;
format!("{sign}{}.{}", &digits[..d], &digits[d..])
}
}
fn cmp_rem_half(r: &str) -> std::cmp::Ordering {
let first = r.as_bytes()[0];
if first > b'5' {
std::cmp::Ordering::Greater
} else if first < b'5' {
std::cmp::Ordering::Less
} else if r[1..].bytes().all(|b| b == b'0') {
std::cmp::Ordering::Equal
} else {
std::cmp::Ordering::Greater
}
}
fn strip_trailing_zeros(d: &str) -> &str {
let end = d.trim_end_matches('0');
if end.is_empty() {
&d[..1]
} else {
end
}
}
fn roundtrips(digits: &str, decpt: i64, x_bits: u64) -> bool {
let e = decpt - digits.len() as i64;
let text = format!("{digits}e{e}");
text.parse::<f64>().map(|v| v.to_bits()) == Ok(x_bits)
}
fn shortest_digits(x: f64, n: &str, decpt: i64) -> (String, i64) {
let x_bits = x.to_bits();
for d in 1..=17usize {
if d >= n.len() {
return (strip_trailing_zeros(n).to_string(), decpt);
}
let lo = &n[..d];
let rem = &n[d..];
let hi_full = inc_decimal(lo);
let (hi, hi_decpt) = if hi_full.len() > d {
(hi_full[..d].to_string(), decpt + 1)
} else {
(hi_full, decpt)
};
let lo_ok = roundtrips(lo, decpt, x_bits);
let hi_ok = roundtrips(&hi, hi_decpt, x_bits);
match (lo_ok, hi_ok) {
(true, false) => return (strip_trailing_zeros(lo).to_string(), decpt),
(false, true) => return (strip_trailing_zeros(&hi).to_string(), hi_decpt),
(true, true) => {
let pick_hi = match cmp_rem_half(rem) {
std::cmp::Ordering::Greater => true,
std::cmp::Ordering::Less => false,
std::cmp::Ordering::Equal => {
(lo.as_bytes()[d - 1] - b'0') % 2 == 1
}
};
return if pick_hi {
(strip_trailing_zeros(&hi).to_string(), hi_decpt)
} else {
(strip_trailing_zeros(lo).to_string(), decpt)
};
}
(false, false) => continue,
}
}
unreachable!("17 significant digits always round-trip a double")
}
fn big_mul_small(v: &mut Vec<u64>, m: u64) {
let mut carry: u128 = 0;
for limb in v.iter_mut() {
let p = (*limb as u128) * (m as u128) + carry;
*limb = p as u64;
carry = p >> 64;
}
while carry > 0 {
v.push(carry as u64);
carry >>= 64;
}
}
fn big_shl(v: &mut Vec<u64>, bits: u64) {
let words = (bits / 64) as usize;
let rem = bits % 64;
if rem > 0 {
let mut carry: u64 = 0;
for limb in v.iter_mut() {
let new = (*limb << rem) | carry;
carry = *limb >> (64 - rem);
*limb = new;
}
if carry > 0 {
v.push(carry);
}
}
if words > 0 {
let mut shifted = vec![0u64; words];
shifted.append(v);
*v = shifted;
}
}
fn big_divmod_small(v: &mut Vec<u64>, d: u64) -> u64 {
let mut rem: u128 = 0;
for limb in v.iter_mut().rev() {
let cur = (rem << 64) | (*limb as u128);
*limb = (cur / d as u128) as u64;
rem = cur % d as u128;
}
while v.len() > 1 && *v.last().unwrap() == 0 {
v.pop();
}
rem as u64
}
fn big_is_zero(v: &[u64]) -> bool {
v.iter().all(|&l| l == 0)
}
fn big_to_decimal(mut v: Vec<u64>) -> String {
const CHUNK: u64 = 10_000_000_000_000_000_000; let mut chunks: Vec<u64> = Vec::new();
loop {
let r = big_divmod_small(&mut v, CHUNK);
chunks.push(r);
if big_is_zero(&v) {
break;
}
}
let mut out = chunks.pop().unwrap().to_string();
for c in chunks.iter().rev() {
out.push_str(&format!("{c:019}"));
}
out
}
fn exact_decimal(x: f64) -> (String, i64) {
debug_assert!(x.is_finite() && x >= 0.0);
let bits = x.to_bits();
let exp_biased = ((bits >> 52) & 0x7ff) as i64;
let frac = bits & ((1u64 << 52) - 1);
let (m, e) = if exp_biased == 0 {
(frac, -1074i64)
} else {
(frac | (1u64 << 52), exp_biased - 1075)
};
if m == 0 {
return ("0".to_string(), 0);
}
let mut v = vec![m];
if e >= 0 {
big_shl(&mut v, e as u64);
(big_to_decimal(v), 0)
} else {
let k = -e;
const POW5_27: u64 = 7_450_580_596_923_828_125; let mut rem = k;
while rem >= 27 {
big_mul_small(&mut v, POW5_27);
rem -= 27;
}
if rem > 0 {
big_mul_small(&mut v, 5u64.pow(rem as u32));
}
(big_to_decimal(v), k)
}
}
fn inc_decimal(q: &str) -> String {
let mut digits: Vec<u8> = q.bytes().collect();
for d in digits.iter_mut().rev() {
if *d == b'9' {
*d = b'0';
} else {
*d += 1;
return String::from_utf8(digits).unwrap();
}
}
let mut out = String::with_capacity(digits.len() + 1);
out.push('1');
out.push_str(std::str::from_utf8(&digits).unwrap());
out
}
fn round_decimal_half_even(n: &str, drop: usize) -> String {
if drop > n.len() {
return "0".to_string();
}
let (q, r) = n.split_at(n.len() - drop);
let q = if q.is_empty() { "0" } else { q };
match cmp_rem_half(r) {
std::cmp::Ordering::Less => q.to_string(),
std::cmp::Ordering::Greater => inc_decimal(q),
std::cmp::Ordering::Equal => {
let last = q.as_bytes()[q.len() - 1];
if (last - b'0') % 2 == 1 {
inc_decimal(q)
} else {
q.to_string()
}
}
}
}
pub fn py_round(x: f64, ndigits: i32) -> f64 {
if !x.is_finite() || x == 0.0 {
return x;
}
let neg = x < 0.0;
let (n, k) = exact_decimal(x.abs());
let nd = ndigits as i64;
if nd >= k {
return x; }
let drop = k - nd;
if drop > n.len() as i64 {
return if neg { -0.0 } else { 0.0 };
}
let q = round_decimal_half_even(&n, drop as usize);
let text = format!("{}{}e{}", if neg { "-" } else { "" }, q, -nd);
text.parse::<f64>().expect("decimal string parses")
}
fn py_fixed(x: f64, nd: usize) -> String {
if x.is_nan() {
return "nan".to_string();
}
if x.is_infinite() {
return if x > 0.0 { "inf" } else { "-inf" }.to_string();
}
let sign = if x.is_sign_negative() { "-" } else { "" };
let (n, k) = exact_decimal(x.abs());
let mut q = if nd as i64 >= k {
let mut s = n;
s.push_str(&"0".repeat((nd as i64 - k) as usize));
s
} else {
let drop = k - nd as i64;
if drop > n.len() as i64 {
"0".to_string()
} else {
round_decimal_half_even(&n, drop as usize)
}
};
if q.len() < nd + 1 {
q = format!("{}{}", "0".repeat(nd + 1 - q.len()), q);
}
if nd == 0 {
format!("{sign}{q}")
} else {
let split = q.len() - nd;
format!("{sign}{}.{}", &q[..split], &q[split..])
}
}
pub fn py_format_1f(x: f64) -> String {
py_fixed(x, 1)
}
pub fn py_format_fixed(x: f64, nd: usize) -> String {
py_fixed(x, nd)
}
pub fn py_format_percent0(x: f64) -> String {
if x.is_nan() {
return "nan%".to_string();
}
if x.is_infinite() {
return if x > 0.0 { "inf%" } else { "-inf%" }.to_string();
}
let mut s = py_fixed(x * 100.0, 0);
s.push('%');
s
}
pub fn py_normpath(path: &str) -> String {
if path.is_empty() {
return ".".to_string();
}
let initial_slashes = if path.starts_with('/') {
if path.starts_with("//") && !path.starts_with("///") {
2
} else {
1
}
} else {
0
};
let mut comps: Vec<&str> = Vec::new();
for comp in path.split('/') {
if comp.is_empty() || comp == "." {
continue;
}
if comp != ".."
|| (initial_slashes == 0 && comps.is_empty())
|| comps.last() == Some(&"..")
{
comps.push(comp);
} else if !comps.is_empty() {
comps.pop();
}
}
let mut out = "/".repeat(initial_slashes);
out.push_str(&comps.join("/"));
if out.is_empty() {
".".to_string()
} else {
out
}
}
pub fn py_abspath(path: &str) -> String {
if path.starts_with('/') {
return py_normpath(path);
}
let cwd = std::env::current_dir()
.map(|p| p.to_string_lossy().into_owned())
.unwrap_or_else(|_| ".".to_string());
py_normpath(&format!("{cwd}/{path}"))
}
pub fn py_relpath(path: &str, start: &str) -> String {
let path_abs = py_abspath(path);
let start_abs = py_abspath(start);
let path_list: Vec<&str> = path_abs.split('/').filter(|c| !c.is_empty()).collect();
let start_list: Vec<&str> = start_abs.split('/').filter(|c| !c.is_empty()).collect();
let common = path_list
.iter()
.zip(start_list.iter())
.take_while(|(a, b)| a == b)
.count();
let mut rel: Vec<&str> = Vec::new();
rel.resize(start_list.len() - common, "..");
rel.extend(&path_list[common..]);
if rel.is_empty() {
".".to_string()
} else {
rel.join("/")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn casefold_basics() {
assert_eq!(py_casefold("Straße"), "strasse");
assert_eq!(py_casefold("ABC"), "abc");
}
#[test]
fn strip_python_whitespace() {
assert_eq!(py_strip("\u{1c}\u{a0} x \t"), "x");
assert_eq!(py_strip("\u{feff}x\u{200b}"), "\u{feff}x\u{200b}");
}
#[test]
fn splitlines_crlf() {
assert_eq!(py_splitlines("a\r\nb\rc\nd\n"), vec!["a", "b", "c", "d"]);
assert_eq!(py_splitlines(""), Vec::<&str>::new());
}
#[test]
fn repr_quote_flip() {
assert_eq!(py_repr_str("it's"), "\"it's\"");
assert_eq!(py_repr_str("both '\""), "'both \\'\"'");
assert_eq!(py_repr_str("café"), "'café'");
assert_eq!(py_repr_str("\u{7f}"), "'\\x7f'");
}
#[test]
fn float_repr_shapes() {
assert_eq!(py_float_repr(1e16), "1e+16");
assert_eq!(py_float_repr(1e-5), "1e-05");
assert_eq!(py_float_repr(0.0001), "0.0001");
assert_eq!(py_float_repr(100.0), "100.0");
assert_eq!(py_float_repr(-0.0), "-0.0");
assert_eq!(py_float_repr(5e-324), "5e-324");
}
#[test]
fn round_half_even_exact() {
assert_eq!(py_round(2.675, 2), 2.67);
assert_eq!(py_round(0.125, 2), 0.12);
assert_eq!(py_round(2.5, 0), 2.0);
assert!(py_round(-0.4, 0) == 0.0 && py_round(-0.4, 0).is_sign_negative());
}
#[test]
fn format_helpers() {
assert_eq!(py_format_1f(0.25), "0.2");
assert_eq!(py_format_1f(-0.04), "-0.0");
assert_eq!(py_format_percent0(0.855), "86%");
}
#[test]
fn stdin_surrogateescape_decode_and_reencode() {
let cases: &[(&[u8], &[u32])] = &[
(b"abc", &[0x61, 0x62, 0x63]),
(b"---\n\xcc\n---", &[0x2d, 0x2d, 0x2d, 0x0a, 0x10FCCC, 0x0a, 0x2d, 0x2d, 0x2d]),
(b"\xc3\x28", &[0x10FCC3, 0x28]), (b"\xf0\x9f\x98", &[0x10FCF0, 0x10FC9F, 0x10FC98]), (b"\xed\xa0\x80", &[0x10FCED, 0x10FCA0, 0x10FC80]), (b"\xc0\xaf", &[0x10FCC0, 0x10FCAF]), (b"\xc3\xa9", &[0xE9]), ];
for (bytes, chars) in cases {
let s = decode_stdin_surrogateescape(bytes);
let got: Vec<u32> = s.chars().map(|c| c as u32).collect();
assert_eq!(&got, chars, "decode of {bytes:?}");
assert_eq!(
encode_stdout_surrogateescape(&s).as_ref(),
*bytes,
"re-encode of {bytes:?}"
);
}
assert!(surrogate_sentinels_active());
}
#[test]
fn sentinel_repr_is_lone_surrogate_escape() {
set_surrogate_sentinels_active(true);
let s = decode_stdin_surrogateescape(b"a\xccb");
assert_eq!(py_repr_str(&s), "'a\\udcccb'"); }
#[test]
fn normpath_contract_examples() {
assert_eq!(py_normpath(""), ".");
assert_eq!(py_normpath("a//b/./c/"), "a/b/c");
assert_eq!(py_normpath("a/b/../c"), "a/c");
assert_eq!(py_normpath("../a"), "../a");
assert_eq!(py_normpath("a/../../b"), "../b");
assert_eq!(py_normpath("/../a"), "/a");
assert_eq!(py_normpath("//a/b"), "//a/b");
assert_eq!(py_normpath("///a/b"), "/a/b");
assert_eq!(py_normpath("/"), "/");
}
#[test]
fn relpath_contract_examples() {
assert_eq!(py_relpath("/x/decisions/decisions/a.md", "/x/decisions"), "decisions/a.md");
assert_eq!(py_relpath("/x/decisions", "/x/decisions"), ".");
assert_eq!(py_relpath("/x/other/a.md", "/x/decisions"), "../other/a.md");
assert_eq!(py_relpath("/x/decisions/a.md", "/x/decisions/"), "a.md");
}
}