diff --git a/crates/capi/src/bytesobject.rs b/crates/capi/src/bytesobject.rs index f4db16af6b3..087bba4a7b3 100644 --- a/crates/capi/src/bytesobject.rs +++ b/crates/capi/src/bytesobject.rs @@ -1,6 +1,6 @@ use crate::object::define_py_check; use crate::{PyObject, pystate::with_vm}; -use core::ffi::c_char; +use core::ffi::{CStr, c_char, c_int}; use rustpython_vm::builtins::PyBytes; define_py_check!(fn PyBytes_Check, types.bytes_type); @@ -29,6 +29,26 @@ pub unsafe extern "C" fn PyBytes_FromStringAndSize( }) } +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyBytes_FromString(s: *const c_char) -> *mut PyObject { + with_vm(|vm| { + let data = unsafe { CStr::from_ptr(s) }.to_bytes().to_vec(); + Ok(vm.ctx.new_bytes(data)) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyBytes_FromObject(obj: *mut PyObject) -> *mut PyObject { + with_vm(|vm| { + let obj = unsafe { &*obj }; + if let Some(bytes) = obj.downcast_ref::() { + Ok(bytes.to_owned()) + } else { + obj.try_bytes_like(vm, |bytes| vm.ctx.new_bytes(bytes.to_vec())) + } + }) +} + #[unsafe(no_mangle)] pub unsafe extern "C" fn PyBytes_Size(bytes: *mut PyObject) -> isize { with_vm(|vm| { @@ -45,6 +65,29 @@ pub unsafe extern "C" fn PyBytes_AsString(bytes: *mut PyObject) -> *mut c_char { }) } +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyBytes_AsStringAndSize( + obj: *mut PyObject, + buffer: *mut *mut c_char, + length: *mut isize, +) -> c_int { + with_vm(|vm| { + let data = unsafe { &*obj }.try_downcast_ref::(vm)?.as_bytes(); + + if length.is_null() { + if data.contains(&0) { + return Err(vm.new_value_error("embedded null byte")); + } + } else { + unsafe { *length = data.len() as isize }; + }; + + unsafe { *buffer = data.as_ptr().cast_mut().cast() }; + + Ok(()) + }) +} + #[cfg(test)] mod tests { use pyo3::prelude::*;