marshal: Use PyAny and PyBytes in safe interface.

This commit is contained in:
Maarten de Vries 2019-04-25 15:05:59 +02:00
parent 35168883b2
commit 739efb1b54

View file

@ -1,7 +1,7 @@
use crate::conversion::AsPyPointer;
use crate::ffi; use crate::ffi;
use crate::{PyErr, PyObject, PyResult, Python}; use crate::types::{PyAny, PyBytes};
use std::os::raw::c_int; use crate::{AsPyPointer, FromPyPointer, PyResult, Python};
use std::os::raw::{c_char, c_int};
/// The current version of the marshal binary format. /// The current version of the marshal binary format.
pub const VERSION: i32 = 4; pub const VERSION: i32 = 4;
@ -10,40 +10,30 @@ pub const VERSION: i32 = 4;
/// ///
/// The built-in marshalling only supports a limited range of object. /// The built-in marshalling only supports a limited range of object.
/// See the [python documentation](https://docs.python.org/3/library/marshal.html) for more details. /// See the [python documentation](https://docs.python.org/3/library/marshal.html) for more details.
pub fn dumps(py: Python, object: &PyObject, version: i32) -> PyResult<Vec<u8>> { pub fn dumps<'a>(py: Python<'a>, object: &impl AsPyPointer, version: i32) -> PyResult<&'a PyBytes> {
let bytes = unsafe { ffi::PyMarshal_WriteObjectToString(object.as_ptr(), version as c_int) };
if bytes.is_null() {
return Err(PyErr::fetch(py));
}
let mut size = 0isize;
let mut data = std::ptr::null_mut();
unsafe { unsafe {
ffi::PyBytes_AsStringAndSize(bytes, &mut data, &mut size); let bytes = ffi::PyMarshal_WriteObjectToString(object.as_ptr(), version as c_int);
let data = Vec::from(std::slice::from_raw_parts(data as *const u8, size as usize)); FromPyPointer::from_owned_ptr_or_err(py, bytes)
ffi::Py_DecRef(bytes);
Ok(data)
} }
} }
/// Deserialize an object from bytes using the Python built-in marshal module. /// Deserialize an object from bytes using the Python built-in marshal module.
pub fn loads(py: Python, data: &impl AsRef<[u8]>) -> PyResult<PyObject> { pub fn loads<'a, B>(py: Python<'a>, data: &B) -> PyResult<&'a PyAny>
where
B: AsRef<[u8]> + ?Sized,
{
let data = data.as_ref(); let data = data.as_ref();
unsafe {
let object = unsafe { let c_str = data.as_ptr() as *const c_char;
ffi::PyMarshal_ReadObjectFromString(data.as_ptr() as *const i8, data.len() as isize) let object = ffi::PyMarshal_ReadObjectFromString(c_str, data.len() as isize);
}; FromPyPointer::from_owned_ptr_or_err(py, object)
if object.is_null() {
return Err(PyErr::fetch(py));
} }
Ok(unsafe { PyObject::from_owned_ptr(py, object) })
} }
#[cfg(test)] #[cfg(test)]
mod test { mod test {
use super::*; use super::*;
use crate::{types::PyDict, ToPyObject}; use crate::types::PyDict;
#[test] #[test]
fn marhshal_roundtrip() { fn marhshal_roundtrip() {
@ -54,15 +44,16 @@ mod test {
dict.set_item("aap", "noot").unwrap(); dict.set_item("aap", "noot").unwrap();
dict.set_item("mies", "wim").unwrap(); dict.set_item("mies", "wim").unwrap();
dict.set_item("zus", "jet").unwrap(); dict.set_item("zus", "jet").unwrap();
let dict = dict.to_object(py);
let bytes = dumps(py, &dict, VERSION).expect("marshalling failed"); let bytes = dumps(py, dict, VERSION)
let deserialzed = loads(py, &bytes).expect("unmarshalling failed"); .expect("marshalling failed")
.as_bytes();
let deserialzed = loads(py, bytes).expect("unmarshalling failed");
assert!(equal(py, &dict, &deserialzed)); assert!(equal(py, dict, deserialzed));
} }
fn equal(_py: Python, a: &PyObject, b: &PyObject) -> bool { fn equal(_py: Python, a: &impl AsPyPointer, b: &impl AsPyPointer) -> bool {
unsafe { ffi::PyObject_RichCompareBool(a.as_ptr(), b.as_ptr(), ffi::Py_EQ) != 0 } unsafe { ffi::PyObject_RichCompareBool(a.as_ptr(), b.as_ptr(), ffi::Py_EQ) != 0 }
} }
} }