Skip to content

Commit

Permalink
Trying doing things the right way, CPython happy, PyPy sad [skip ci]
Browse files Browse the repository at this point in the history
  • Loading branch information
@milesgranger
milesgranger committed Mar 1, 2024
1 parent a477ed3 commit 51e0988
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 57 deletions.
65 changes: 20 additions & 45 deletions cramjam-python/src/io.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ use std::path::PathBuf;

pub(crate) trait AsBytes {
fn as_bytes(&self) -> &[u8];
fn as_bytes_mut(&mut self) -> &mut [u8];
fn as_bytes_mut(&mut self) -> PyResult<&mut [u8]>;
}

/// A native Rust file-like object. Reading and writing takes place
Expand Down Expand Up @@ -49,7 +49,7 @@ impl AsBytes for RustyFile {
entire file into memory; consider using cramjam.Buffer"
)
}
fn as_bytes_mut(&mut self) -> &mut [u8] {
fn as_bytes_mut(&mut self) -> PyResult<&mut [u8]> {
unimplemented!(
"Converting a File to bytes is not supported, as it'd require reading the \
entire file into memory; consider using cramjam.Buffer"
Expand Down Expand Up @@ -194,21 +194,24 @@ impl PythonBuffer {
}
/// Is the Python buffer readonly
pub fn readonly(&self) -> bool {
self.inner.readonly != 0
self.inner.readonly == 1
}
/// Get the underlying buffer as a slice of bytes
pub fn as_slice(&self) -> &[u8] {
unsafe { std::slice::from_raw_parts(self.buf_ptr() as *const u8, self.len_bytes()) }
}
/// Get the underlying buffer as a mutable slice of bytes
pub fn as_slice_mut(&mut self) -> PyResult<&mut [u8]> {
// TODO: For v3 release, add self.readonly check; bytes is readonly but
// v1 and v2 releases have not treated it as such.
if self.readonly() {
let repr = Python::with_gil(|py| unsafe { PyObject::from_borrowed_ptr(py, self.inner.obj) }.to_string());
let msg = format!("The output buffer '{}' is readonly, refusing to overwrite.", repr);
return Err(pyo3::exceptions::PyTypeError::new_err(msg));
}
Ok(unsafe { std::slice::from_raw_parts_mut(self.buf_ptr() as *mut u8, self.len_bytes()) })
}
/// If underlying buffer is c_contiguous
pub fn is_c_contiguous(&self) -> bool {
unsafe { ffi::PyBuffer_IsContiguous(&*self.inner as *const ffi::Py_buffer, b'C' as std::os::raw::c_char) != 0 }
unsafe { ffi::PyBuffer_IsContiguous(&*self.inner as *const ffi::Py_buffer, b'C' as std::os::raw::c_char) == 1 }
}
/// Dimensions for buffer
pub fn dimensions(&self) -> usize {
Expand Down Expand Up @@ -244,46 +247,17 @@ impl<'py> FromPyObject<'py> for PythonBuffer {
}
}

#[cfg(not(PyPy))]
fn make_py_buffer(obj: &PyAny) -> PyResult<ffi::Py_buffer> {
let mut buf = mem::MaybeUninit::uninit();
let rc = unsafe { ffi::PyObject_GetBuffer(obj.as_ptr(), buf.as_mut_ptr(), ffi::PyBUF_CONTIG_RO) };
if rc != 0 {
return Err(exceptions::PyBufferError::new_err(
"Failed to get buffer, is it C contiguous, and shape is not null?",
));
}
let buf = unsafe { mem::MaybeUninit::<ffi::Py_buffer>::assume_init(buf) };
Ok(buf)
}

#[cfg(PyPy)]
fn make_py_buffer(obj: &PyAny) -> PyResult<ffi::Py_buffer> {
let is_memview = unsafe { ffi::PyMemoryView_Check(obj.as_ptr()) } == 1;

let mut object = Python::with_gil(|py| obj.to_object(py));

if !is_memview {
let ptr = unsafe { ffi::PyMemoryView_FromObject(obj.as_ptr()) };
Python::with_gil(|py| {
object = unsafe { PyObject::from_owned_ptr(py, ptr) };
})
}
let mut buf = mem::MaybeUninit::uninit();
let rc = unsafe { ffi::PyObject_GetBuffer(object.as_ptr(), buf.as_mut_ptr(), ffi::PyBUF_CONTIG_RO) };
if rc != 0 {
return Err(exceptions::PyBufferError::new_err(
"Failed to get buffer, is it C contiguous, and shape is not null?",
));
}
let buf = unsafe { mem::MaybeUninit::<ffi::Py_buffer>::assume_init(buf) };
Ok(buf)
}

impl<'py> TryFrom<&'py PyAny> for PythonBuffer {
type Error = PyErr;
fn try_from(obj: &'py PyAny) -> Result<Self, Self::Error> {
let py_buffer = make_py_buffer(obj)?;
let mut buf = mem::MaybeUninit::uninit();
let rc = unsafe { ffi::PyObject_GetBuffer(obj.as_ptr(), buf.as_mut_ptr(), ffi::PyBUF_CONTIG_RO) };
if rc != 0 {
return Err(exceptions::PyBufferError::new_err(
"Failed to get buffer, is it C contiguous, and shape is not null?",
));
}
let py_buffer = unsafe { mem::MaybeUninit::<ffi::Py_buffer>::assume_init(buf) };
let buf = Self {
inner: std::pin::Pin::from(Box::new(py_buffer)),
pos: 0,
Expand Down Expand Up @@ -357,8 +331,9 @@ impl AsBytes for RustyBuffer {
fn as_bytes(&self) -> &[u8] {
self.inner.get_ref().as_slice()
}
fn as_bytes_mut(&mut self) -> &mut [u8] {
self.inner.get_mut().as_mut_slice()
fn as_bytes_mut(&mut self) -> PyResult<&mut [u8]> {
let slice = self.inner.get_mut().as_mut_slice();
Ok(slice)
}
}

Expand Down
16 changes: 8 additions & 8 deletions cramjam-python/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,18 +101,18 @@ impl<'a> AsBytes for BytesType<'a> {
}
}
}
fn as_bytes_mut(&mut self) -> &mut [u8] {
fn as_bytes_mut(&mut self) -> PyResult<&mut [u8]> {
match self {
BytesType::RustyBuffer(b) => {
let mut py_ref = b.borrow_mut();
let bytes = py_ref.as_bytes_mut();
unsafe { std::slice::from_raw_parts_mut(bytes.as_mut_ptr(), bytes.len()) }
let bytes = py_ref.as_bytes_mut()?;
Ok(unsafe { std::slice::from_raw_parts_mut(bytes.as_mut_ptr(), bytes.len()) })
}
BytesType::PyBuffer(b) => b.as_slice_mut().unwrap(),
BytesType::PyBuffer(b) => b.as_slice_mut(),
BytesType::RustyFile(b) => {
let mut py_ref = b.borrow_mut();
let bytes = py_ref.as_bytes_mut();
unsafe { std::slice::from_raw_parts_mut(bytes.as_mut_ptr(), bytes.len()) }
let bytes = py_ref.as_bytes_mut()?;
Ok(unsafe { std::slice::from_raw_parts_mut(bytes.as_mut_ptr(), bytes.len()) })
}
}
}
Expand Down Expand Up @@ -212,7 +212,7 @@ macro_rules! generic {
})
},
_ => {
let bytes_out = $output.as_bytes_mut();
let bytes_out = $output.as_bytes_mut()?;
$py.allow_threads(|| {
$op(f_in, &mut Cursor::new(bytes_out) $(, $level)? )
})
Expand All @@ -237,7 +237,7 @@ macro_rules! generic {
})
},
_ => {
let bytes_out = $output.as_bytes_mut();
let bytes_out = $output.as_bytes_mut()?;
$py.allow_threads(|| {
$op(bytes_in, &mut Cursor::new(bytes_out) $(, $level)?)
})
Expand Down
4 changes: 2 additions & 2 deletions cramjam-python/src/lz4.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ pub fn compress_block(
#[pyfunction]
pub fn decompress_block_into(py: Python, input: BytesType, mut output: BytesType) -> PyResult<usize> {
let bytes = input.as_bytes();
let out_bytes = output.as_bytes_mut();
let out_bytes = output.as_bytes_mut()?;
py.allow_threads(|| libcramjam::lz4::block::decompress_into(bytes, out_bytes, Some(true)))
.map_err(DecompressionError::from_err)
.map(|v| v as _)
Expand Down Expand Up @@ -180,7 +180,7 @@ pub fn compress_block_into(
store_size: Option<bool>,
) -> PyResult<usize> {
let bytes = data.as_bytes();
let out_bytes = output.as_bytes_mut();
let out_bytes = output.as_bytes_mut()?;
py.allow_threads(|| {
libcramjam::lz4::block::compress_into(bytes, out_bytes, compression.map(|v| v as _), acceleration, store_size)
})
Expand Down
4 changes: 2 additions & 2 deletions cramjam-python/src/snappy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ pub fn decompress_into(py: Python, input: BytesType, mut output: BytesType) -> P
#[pyfunction]
pub fn compress_raw_into(py: Python, input: BytesType, mut output: BytesType) -> PyResult<usize> {
let bytes_in = input.as_bytes();
let bytes_out = output.as_bytes_mut();
let bytes_out = output.as_bytes_mut()?;
py.allow_threads(|| libcramjam::snappy::raw::compress(bytes_in, bytes_out))
.map_err(CompressionError::from_err)
}
Expand All @@ -109,7 +109,7 @@ pub fn compress_raw_into(py: Python, input: BytesType, mut output: BytesType) ->
#[pyfunction]
pub fn decompress_raw_into(py: Python, input: BytesType, mut output: BytesType) -> PyResult<usize> {
let bytes_in = input.as_bytes();
let bytes_out = output.as_bytes_mut();
let bytes_out = output.as_bytes_mut()?;
py.allow_threads(|| libcramjam::snappy::raw::decompress(bytes_in, bytes_out))
.map_err(DecompressionError::from_err)
}
Expand Down

0 comments on commit 51e0988

Please sign in to comment.