Skip to content

Commit 72daf7c

Browse files
Add more unicode functions to the c-api (#8272)
* Add more unicode functions Review Add unicode decode functions Add more functions * Remove * Add helper
1 parent ea79826 commit 72daf7c

2 files changed

Lines changed: 321 additions & 21 deletions

File tree

crates/capi/src/unicodeobject.rs

Lines changed: 317 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@ use core::ffi::{CStr, c_char, c_int};
55
use core::ptr::NonNull;
66
use core::slice;
77
use core::str;
8-
use rustpython_vm::builtins::{PyStr, PyStrRef};
9-
use rustpython_vm::{PyObjectRef, PyResult, VirtualMachine};
8+
use rustpython_vm::builtins::{PyBytesRef, PyStr, PyStrRef, PyUtf8StrRef};
9+
use rustpython_vm::common::wtf8::{CodePoint, Wtf8Buf};
10+
use rustpython_vm::convert::ToPyObject;
11+
use rustpython_vm::{AsObject, PyObjectRef, PyResult, VirtualMachine};
1012

1113
define_py_check!(fn PyUnicode_Check, types.str_type);
1214
define_py_check!(exact fn PyUnicode_CheckExact, types.str_type);
@@ -37,6 +39,36 @@ pub unsafe extern "C" fn PyUnicode_FromStringAndSize(
3739
})
3840
}
3941

42+
#[unsafe(no_mangle)]
43+
pub unsafe extern "C" fn PyUnicode_FromString(s: *const c_char) -> *mut PyObject {
44+
with_vm(|vm| {
45+
let s = unsafe { s.try_as_str(vm)? };
46+
Ok(vm.ctx.new_str(s))
47+
})
48+
}
49+
50+
#[unsafe(no_mangle)]
51+
pub unsafe extern "C" fn PyUnicode_FromObject(obj: *mut PyObject) -> *mut PyObject {
52+
with_vm(|vm| {
53+
Ok(unsafe { &*obj }
54+
.try_downcast_ref::<PyStr>(vm)?
55+
.as_object()
56+
.str(vm))
57+
})
58+
}
59+
60+
#[unsafe(no_mangle)]
61+
pub unsafe extern "C" fn PyUnicode_FromOrdinal(ordinal: c_int) -> *mut PyObject {
62+
with_vm(|vm| {
63+
let ordinal: u32 = ordinal
64+
.try_into()
65+
.map_err(|_| vm.new_value_error("ordinal not in range(0x110000)"))?;
66+
let code_point = CodePoint::from_u32(ordinal)
67+
.ok_or_else(|| vm.new_value_error("ordinal not in range(0x110000)"))?;
68+
Ok(vm.ctx.new_str(Wtf8Buf::from_iter([code_point])))
69+
})
70+
}
71+
4072
#[unsafe(no_mangle)]
4173
pub unsafe extern "C" fn PyUnicode_AsUTF8AndSize(
4274
obj: *mut PyObject,
@@ -61,37 +93,155 @@ pub unsafe extern "C" fn PyUnicode_AsUTF8AndSize(
6193
})
6294
}
6395

96+
fn encode_unicode(
97+
vm: &VirtualMachine,
98+
unicode: *mut PyObject,
99+
encoding: &str,
100+
errors: Option<PyUtf8StrRef>,
101+
) -> PyResult<PyBytesRef> {
102+
let unicode = unsafe { &*unicode }
103+
.try_downcast_ref::<PyStr>(vm)?
104+
.to_owned();
105+
vm.state
106+
.codec_registry
107+
.encode_text(unicode, encoding, errors, vm)
108+
}
109+
110+
#[unsafe(no_mangle)]
111+
pub unsafe extern "C" fn PyUnicode_AsASCIIString(unicode: *mut PyObject) -> *mut PyObject {
112+
with_vm(|vm| encode_unicode(vm, unicode, "ascii", None))
113+
}
114+
115+
#[unsafe(no_mangle)]
116+
pub unsafe extern "C" fn PyUnicode_AsLatin1String(unicode: *mut PyObject) -> *mut PyObject {
117+
with_vm(|vm| encode_unicode(vm, unicode, "latin-1", None))
118+
}
119+
120+
#[unsafe(no_mangle)]
121+
pub unsafe extern "C" fn PyUnicode_AsRawUnicodeEscapeString(
122+
unicode: *mut PyObject,
123+
) -> *mut PyObject {
124+
with_vm(|vm| encode_unicode(vm, unicode, "raw-unicode-escape", None))
125+
}
126+
127+
#[unsafe(no_mangle)]
128+
pub unsafe extern "C" fn PyUnicode_AsUTF16String(unicode: *mut PyObject) -> *mut PyObject {
129+
with_vm(|vm| encode_unicode(vm, unicode, "utf-16", None))
130+
}
131+
132+
#[unsafe(no_mangle)]
133+
pub unsafe extern "C" fn PyUnicode_AsUTF32String(unicode: *mut PyObject) -> *mut PyObject {
134+
with_vm(|vm| encode_unicode(vm, unicode, "utf-32", None))
135+
}
136+
137+
#[unsafe(no_mangle)]
138+
pub unsafe extern "C" fn PyUnicode_AsUnicodeEscapeString(unicode: *mut PyObject) -> *mut PyObject {
139+
with_vm(|vm| encode_unicode(vm, unicode, "unicode-escape", None))
140+
}
141+
64142
#[unsafe(no_mangle)]
65143
pub unsafe extern "C" fn PyUnicode_AsEncodedString(
66144
unicode: *mut PyObject,
67145
encoding: *const c_char,
68146
errors: *const c_char,
69147
) -> *mut PyObject {
70148
with_vm(|vm| {
71-
let unicode = unsafe { &*unicode }
72-
.try_downcast_ref::<PyStr>(vm)?
73-
.to_owned();
74149
let encoding = unsafe { encoding.try_as_str_opt(vm) }?.unwrap_or("utf-8");
75150
let errors =
76151
unsafe { errors.try_as_str_opt(vm) }?.map(|errors| vm.ctx.new_utf8_str(errors));
77-
vm.state
78-
.codec_registry
79-
.encode_text(unicode, encoding, errors, vm)
152+
encode_unicode(vm, unicode, encoding, errors)
80153
})
81154
}
82155

83156
#[unsafe(no_mangle)]
84157
pub unsafe extern "C" fn PyUnicode_AsUTF8String(unicode: *mut PyObject) -> *mut PyObject {
158+
with_vm(|vm| encode_unicode(vm, unicode, "utf-8", None))
159+
}
160+
161+
#[unsafe(no_mangle)]
162+
pub unsafe extern "C" fn PyUnicode_Decode(
163+
s: *const c_char,
164+
size: isize,
165+
encoding: *const c_char,
166+
errors: *const c_char,
167+
) -> *mut PyObject {
85168
with_vm(|vm| {
86-
let unicode = unsafe { &*unicode }
87-
.try_downcast_ref::<PyStr>(vm)?
88-
.to_owned();
169+
let size: usize = size
170+
.try_into()
171+
.map_err(|_| vm.new_system_error("size must be non-negative"))?;
172+
173+
let bytes = if s.is_null() {
174+
if size != 0 {
175+
return Err(vm.new_system_error("decode called with null data and non-zero size"));
176+
}
177+
Vec::new()
178+
} else {
179+
unsafe { slice::from_raw_parts(s.cast::<u8>(), size) }.to_vec()
180+
};
181+
182+
let encoding = unsafe { encoding.try_as_str_opt(vm)?.unwrap_or("utf-8") };
183+
let errors =
184+
unsafe { errors.try_as_str_opt(vm) }?.map(|errors| vm.ctx.new_utf8_str(errors));
185+
89186
vm.state
90187
.codec_registry
91-
.encode_text(unicode, "utf-8", None, vm)
188+
.decode_text(vm.ctx.new_bytes(bytes).into(), encoding, errors, vm)
92189
})
93190
}
94191

192+
#[unsafe(no_mangle)]
193+
pub unsafe extern "C" fn PyUnicode_DecodeASCII(
194+
s: *const c_char,
195+
size: isize,
196+
errors: *const c_char,
197+
) -> *mut PyObject {
198+
unsafe { PyUnicode_Decode(s, size, c"ascii".as_ptr(), errors) }
199+
}
200+
201+
#[unsafe(no_mangle)]
202+
pub unsafe extern "C" fn PyUnicode_DecodeLatin1(
203+
s: *const c_char,
204+
size: isize,
205+
errors: *const c_char,
206+
) -> *mut PyObject {
207+
unsafe { PyUnicode_Decode(s, size, c"latin-1".as_ptr(), errors) }
208+
}
209+
210+
#[unsafe(no_mangle)]
211+
pub unsafe extern "C" fn PyUnicode_DecodeRawUnicodeEscape(
212+
s: *const c_char,
213+
size: isize,
214+
errors: *const c_char,
215+
) -> *mut PyObject {
216+
unsafe { PyUnicode_Decode(s, size, c"raw-unicode-escape".as_ptr(), errors) }
217+
}
218+
219+
#[unsafe(no_mangle)]
220+
pub unsafe extern "C" fn PyUnicode_DecodeUTF7(
221+
s: *const c_char,
222+
size: isize,
223+
errors: *const c_char,
224+
) -> *mut PyObject {
225+
unsafe { PyUnicode_Decode(s, size, c"utf-7".as_ptr(), errors) }
226+
}
227+
228+
#[unsafe(no_mangle)]
229+
pub unsafe extern "C" fn PyUnicode_DecodeUTF8(
230+
s: *const c_char,
231+
size: isize,
232+
errors: *const c_char,
233+
) -> *mut PyObject {
234+
unsafe { PyUnicode_Decode(s, size, c"utf-8".as_ptr(), errors) }
235+
}
236+
#[unsafe(no_mangle)]
237+
pub unsafe extern "C" fn PyUnicode_DecodeUnicodeEscape(
238+
s: *const c_char,
239+
size: isize,
240+
errors: *const c_char,
241+
) -> *mut PyObject {
242+
unsafe { PyUnicode_Decode(s, size, c"unicode-escape".as_ptr(), errors) }
243+
}
244+
95245
#[unsafe(no_mangle)]
96246
pub unsafe extern "C" fn PyUnicode_DecodeFSDefaultAndSize(
97247
s: *const c_char,
@@ -106,6 +256,89 @@ pub unsafe extern "C" fn PyUnicode_DecodeFSDefaultAndSize(
106256
})
107257
}
108258

259+
#[unsafe(no_mangle)]
260+
pub unsafe extern "C" fn PyUnicode_Concat(
261+
left: *mut PyObject,
262+
right: *mut PyObject,
263+
) -> *mut PyObject {
264+
with_vm(|vm| {
265+
let left = unsafe { &*left }.try_downcast_ref::<PyStr>(vm)?;
266+
let right = unsafe { &*right }.try_downcast_ref::<PyStr>(vm)?;
267+
vm._add(left.as_object(), right.as_object())
268+
})
269+
}
270+
271+
#[unsafe(no_mangle)]
272+
pub unsafe extern "C" fn PyUnicode_GetLength(unicode: *mut PyObject) -> isize {
273+
with_vm(|vm| {
274+
let unicode = unsafe { &*unicode }.try_downcast_ref::<PyStr>(vm)?;
275+
Ok(unicode.char_len())
276+
})
277+
}
278+
279+
#[unsafe(no_mangle)]
280+
pub unsafe extern "C" fn PyUnicode_GetDefaultEncoding() -> *const c_char {
281+
c"utf-8".as_ptr()
282+
}
283+
284+
#[unsafe(no_mangle)]
285+
pub unsafe extern "C" fn PyUnicode_InternFromString(s: *const c_char) -> *mut PyObject {
286+
with_vm(|vm| {
287+
let s = unsafe { s.try_as_str(vm)? };
288+
Ok(vm.ctx.intern_str(s).to_owned())
289+
})
290+
}
291+
292+
#[unsafe(no_mangle)]
293+
pub unsafe extern "C" fn PyUnicode_Compare(left: *mut PyObject, right: *mut PyObject) -> c_int {
294+
with_vm(|vm| {
295+
let left = unsafe { &*left }.try_downcast_ref::<PyStr>(vm)?;
296+
let right = unsafe { &*right }.try_downcast_ref::<PyStr>(vm)?;
297+
Ok(match left.as_wtf8().cmp(right.as_wtf8()) {
298+
core::cmp::Ordering::Less => -1,
299+
core::cmp::Ordering::Equal => 0,
300+
core::cmp::Ordering::Greater => 1,
301+
})
302+
})
303+
}
304+
305+
#[unsafe(no_mangle)]
306+
pub unsafe extern "C" fn PyUnicode_CompareWithASCIIString(
307+
left: *mut PyObject,
308+
right: *const c_char,
309+
) -> c_int {
310+
with_vm(|vm| {
311+
let left = unsafe { &*left }.try_downcast_ref::<PyStr>(vm)?;
312+
let right = unsafe { right.try_as_str(vm)? };
313+
Ok(match left.as_wtf8().cmp(right.into()) {
314+
core::cmp::Ordering::Less => -1,
315+
core::cmp::Ordering::Equal => 0,
316+
core::cmp::Ordering::Greater => 1,
317+
})
318+
})
319+
}
320+
321+
#[unsafe(no_mangle)]
322+
pub unsafe extern "C" fn PyUnicode_Equal(left: *mut PyObject, right: *mut PyObject) -> c_int {
323+
with_vm(|vm| {
324+
let left = unsafe { &*left }.try_downcast_ref::<PyStr>(vm)?;
325+
let right = unsafe { &*right }.try_downcast_ref::<PyStr>(vm)?;
326+
Ok(left.as_wtf8() == right.as_wtf8())
327+
})
328+
}
329+
330+
#[unsafe(no_mangle)]
331+
pub unsafe extern "C" fn PyUnicode_EqualToUTF8(
332+
unicode: *mut PyObject,
333+
string: *const c_char,
334+
) -> c_int {
335+
with_vm(|vm| {
336+
let unicode = unsafe { &*unicode }.try_downcast_ref::<PyStr>(vm)?;
337+
let other = unsafe { string.try_as_str(vm)? };
338+
Ok(unicode.to_str().is_some_and(|s| s == other))
339+
})
340+
}
341+
109342
#[unsafe(no_mangle)]
110343
pub unsafe extern "C" fn PyUnicode_DecodeFSDefault(s: *const c_char) -> *mut PyObject {
111344
with_vm(|vm| {
@@ -141,14 +374,11 @@ pub(crate) fn decode_fsdefault_and_size(
141374
#[unsafe(no_mangle)]
142375
pub unsafe extern "C" fn PyUnicode_EncodeFSDefault(unicode: *mut PyObject) -> *mut PyObject {
143376
with_vm(|vm| {
144-
let unicode = unsafe { &*unicode }
145-
.try_downcast_ref::<PyStr>(vm)?
146-
.to_owned();
147-
vm.state.codec_registry.encode_text(
377+
encode_unicode(
378+
vm,
148379
unicode,
149380
vm.fs_encoding().as_str(),
150381
Some(vm.fs_encode_errors().to_owned()),
151-
vm,
152382
)
153383
})
154384
}
@@ -181,6 +411,76 @@ pub unsafe extern "C" fn PyUnicode_FromEncodedObject(
181411
})
182412
}
183413

414+
#[unsafe(no_mangle)]
415+
pub unsafe extern "C" fn PyUnicode_Contains(
416+
container: *mut PyObject,
417+
element: *mut PyObject,
418+
) -> c_int {
419+
with_vm(|vm| {
420+
let container = unsafe { &*container }.try_downcast_ref::<PyStr>(vm)?;
421+
let element = unsafe { &*element }.try_downcast_ref::<PyStr>(vm)?;
422+
Ok(container.as_wtf8().contains(element.as_wtf8()))
423+
})
424+
}
425+
426+
#[unsafe(no_mangle)]
427+
pub unsafe extern "C" fn PyUnicode_Format(
428+
format: *mut PyObject,
429+
args: *mut PyObject,
430+
) -> *mut PyObject {
431+
with_vm(|vm| {
432+
let format = unsafe { &*format }.try_downcast_ref::<PyStr>(vm)?;
433+
let result = format.__mod__(unsafe { &*args }.to_owned(), vm)?;
434+
Ok(result.to_pyobject(vm))
435+
})
436+
}
437+
438+
#[unsafe(no_mangle)]
439+
pub unsafe extern "C" fn PyUnicode_IsIdentifier(s: *mut PyObject) -> c_int {
440+
with_vm(|vm| {
441+
let s = unsafe { &*s }.try_downcast_ref::<PyStr>(vm)?;
442+
Ok(s.isidentifier())
443+
})
444+
}
445+
446+
#[unsafe(no_mangle)]
447+
pub unsafe extern "C" fn PyUnicode_Partition(
448+
s: *mut PyObject,
449+
sep: *mut PyObject,
450+
) -> *mut PyObject {
451+
with_vm(|vm| {
452+
let s = unsafe { &*s }.try_downcast_ref::<PyStr>(vm)?;
453+
let sep = unsafe { &*sep }.try_downcast_ref::<PyStr>(vm)?;
454+
s.partition(sep.to_owned(), vm)
455+
})
456+
}
457+
458+
#[unsafe(no_mangle)]
459+
pub unsafe extern "C" fn PyUnicode_RPartition(
460+
s: *mut PyObject,
461+
sep: *mut PyObject,
462+
) -> *mut PyObject {
463+
with_vm(|vm| {
464+
let s = unsafe { &*s }.try_downcast_ref::<PyStr>(vm)?;
465+
let sep = unsafe { &*sep }.try_downcast_ref::<PyStr>(vm)?;
466+
s.rpartition(sep.to_owned(), vm)
467+
})
468+
}
469+
470+
#[unsafe(no_mangle)]
471+
pub unsafe extern "C" fn PyUnicode_Translate(
472+
str_obj: *mut PyObject,
473+
table: *mut PyObject,
474+
_errors: *const c_char,
475+
) -> *mut PyObject {
476+
with_vm(|vm| {
477+
let str_obj = unsafe { &*str_obj }.try_downcast_ref::<PyStr>(vm)?;
478+
Ok(str_obj
479+
.translate(unsafe { &*table }.to_owned(), vm)?
480+
.to_pyobject(vm))
481+
})
482+
}
483+
184484
#[unsafe(no_mangle)]
185485
pub unsafe extern "C" fn PyUnicode_InternInPlace(string: *mut *mut PyObject) {
186486
with_vm(|vm| {

0 commit comments

Comments
 (0)