Skip to content

Commit dc84268

Browse files
authored
Align Linux locale collation with Python 3.15 (#9674)
Use wide collation strings, retain surrogate code points, encode transform weights in Python 3.15 order, and preserve errno and allocation failures. AI-assisted implementation, source review, and validation using OpenAI Codex. Assisted-by: OpenAI Codex:model version not exposed
1 parent 36e100b commit dc84268

2 files changed

Lines changed: 130 additions & 2 deletions

File tree

‎crates/host_env/src/locale.rs‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,103 @@ pub fn strxfrm(string: &CStr, _initial_len: usize) -> Vec<u8> {
222222
buff
223223
}
224224

225+
#[cfg(target_os = "linux")]
226+
pub use wide_collation::{CollationString, CollationTransformError};
227+
228+
#[cfg(target_os = "linux")]
229+
mod wide_collation {
230+
use rustpython_wtf8::{CodePoint, Wtf8, Wtf8Buf};
231+
use std::io;
232+
233+
unsafe extern "C" {
234+
fn wcscoll(first: *const libc::wchar_t, second: *const libc::wchar_t) -> libc::c_int;
235+
fn wcsxfrm(
236+
destination: *mut libc::wchar_t,
237+
source: *const libc::wchar_t,
238+
capacity: usize,
239+
) -> usize;
240+
}
241+
242+
/// A NUL-terminated wide string for the locale collation APIs.
243+
pub struct CollationString(Vec<libc::wchar_t>);
244+
245+
pub enum CollationTransformError {
246+
Memory,
247+
Os(io::Error),
248+
}
249+
250+
impl CollationString {
251+
/// Preserve Python's lone surrogates and reject embedded NULs.
252+
pub fn new(string: &Wtf8) -> Option<Self> {
253+
let mut wide = Vec::with_capacity(string.len() + 1);
254+
for point in string.code_points() {
255+
let value = point.to_u32();
256+
if value == 0 {
257+
return None;
258+
}
259+
wide.push(value as libc::wchar_t);
260+
}
261+
wide.push(0);
262+
Some(Self(wide))
263+
}
264+
265+
pub fn compare(&self, other: &Self) -> libc::c_int {
266+
// Both buffers retain their terminator and cannot contain interior NULs.
267+
unsafe { wcscoll(self.0.as_ptr(), other.0.as_ptr()) }
268+
}
269+
270+
pub fn transform(&self) -> Result<Wtf8Buf, CollationTransformError> {
271+
let mut dummy = [0];
272+
crate::os::clear_errno();
273+
let needed = unsafe { wcsxfrm(dummy.as_mut_ptr(), self.0.as_ptr(), 1) };
274+
let errno = crate::os::get_errno();
275+
if errno != 0 && errno != libc::ERANGE {
276+
return Err(CollationTransformError::Os(io::Error::from_raw_os_error(
277+
errno,
278+
)));
279+
}
280+
let capacity = needed
281+
.checked_add(1)
282+
.ok_or(CollationTransformError::Memory)?;
283+
let mut wide = Vec::new();
284+
wide.try_reserve_exact(capacity)
285+
.map_err(|_| CollationTransformError::Memory)?;
286+
wide.resize(capacity, 0);
287+
crate::os::clear_errno();
288+
let written = unsafe { wcsxfrm(wide.as_mut_ptr(), self.0.as_ptr(), wide.len()) };
289+
let errno = crate::os::get_errno();
290+
if errno != 0 {
291+
return Err(CollationTransformError::Os(io::Error::from_raw_os_error(
292+
errno,
293+
)));
294+
}
295+
// A concurrent locale change can invalidate the size query. Never
296+
// interpret a truncated or unterminated result as a collation key.
297+
if written >= wide.len() {
298+
return Err(CollationTransformError::Os(io::Error::from_raw_os_error(
299+
libc::ERANGE,
300+
)));
301+
}
302+
wide.truncate(written);
303+
304+
// wcsxfrm emits integer weights, not necessarily Unicode. CPython
305+
// 3.15 splits weights above 0x10000 into two ordered code points.
306+
// Keep individual surrogates rather than decoding as UTF-16.
307+
let mut key = Wtf8Buf::new();
308+
for weight in wide {
309+
let value = weight as u32;
310+
if value > 0x10000 {
311+
key.push(CodePoint::from_u32((value >> 16) | 0x10000).unwrap());
312+
key.push(CodePoint::from_u32(value & 0xffff).unwrap());
313+
} else {
314+
key.push(CodePoint::from_u32(value).unwrap());
315+
}
316+
}
317+
Ok(key)
318+
}
319+
}
320+
}
321+
225322
pub fn setlocale(category: i32, locale: Option<&CStr>) -> Option<Vec<u8>> {
226323
let result = unsafe {
227324
match locale {

‎crates/stdlib/src/locale.rs‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ pub(crate) use _locale::module_def;
66
mod _locale {
77
use alloc::ffi::CString;
88
use rustpython_host_env::locale as host_locale;
9+
#[cfg(target_os = "linux")]
10+
use rustpython_vm::builtins::{PyStr, PyStrRef};
911
use rustpython_vm::{
1012
PyObjectRef, PyResult, VirtualMachine,
1113
builtins::{PyDictRef, PyIntRef, PyListRef, PyTypeRef, PyUtf8StrRef},
@@ -90,7 +92,25 @@ mod _locale {
9092
Ok(vm.new_pyobj(host_locale::wcscoll(&w1, &w2)))
9193
}
9294

93-
#[cfg(not(windows))]
95+
#[cfg(target_os = "linux")]
96+
fn collation_string(
97+
string: &PyStr,
98+
vm: &VirtualMachine,
99+
) -> PyResult<host_locale::CollationString> {
100+
host_locale::CollationString::new(string.as_wtf8())
101+
.ok_or_else(|| vm.new_value_error("embedded null character"))
102+
}
103+
104+
#[cfg(target_os = "linux")]
105+
#[pyfunction]
106+
fn strcoll(os1: PyStrRef, os2: PyStrRef, vm: &VirtualMachine) -> PyResult {
107+
// Bind both strings before checking for NULs, as Argument Clinic does.
108+
let first = collation_string(&os1, vm)?;
109+
let second = collation_string(&os2, vm)?;
110+
Ok(vm.new_pyobj(first.compare(&second)))
111+
}
112+
113+
#[cfg(not(any(windows, target_os = "linux")))]
94114
#[pyfunction]
95115
fn strcoll(os1: PyUtf8StrRef, os2: PyUtf8StrRef, vm: &VirtualMachine) -> PyResult {
96116
let cstr1 = CString::new(os1.as_str()).map_err(|e| e.to_pyexception(vm))?;
@@ -109,7 +129,18 @@ mod _locale {
109129
Ok(vm.new_pyobj(String::from_utf16_lossy(&transformed)))
110130
}
111131

112-
#[cfg(not(windows))]
132+
#[cfg(target_os = "linux")]
133+
#[pyfunction]
134+
fn strxfrm(string: PyStrRef, vm: &VirtualMachine) -> PyResult {
135+
let string = collation_string(&string, vm)?;
136+
let key = string.transform().map_err(|err| match err {
137+
host_locale::CollationTransformError::Memory => vm.new_memory_error(""),
138+
host_locale::CollationTransformError::Os(err) => err.to_pyexception(vm),
139+
})?;
140+
Ok(vm.ctx.new_str(key).into())
141+
}
142+
143+
#[cfg(not(any(windows, target_os = "linux")))]
113144
#[pyfunction]
114145
fn strxfrm(string: PyUtf8StrRef, vm: &VirtualMachine) -> PyResult {
115146
// https://github.com/python/cpython/blob/eaae563b6878aa050b4ad406b67728b6b066220e/Modules/_localemodule.c#L390-L442

0 commit comments

Comments
 (0)