Skip to content

Instantly share code, notes, and snippets.

@cswinter
Created Jan 21, 2021
Embed
What would you like to do?
PyArcVec
mod arc_buffer {
use lazy_static::lazy_static;
use pyo3::ffi;
use pyo3::prelude::*;
use pyo3::AsPyPointer;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
lazy_static! {
static ref COUNTER: AtomicUsize = AtomicUsize::default();
}
#[pyclass]
pub struct ArcVec {
data: Arc<Vec<u8>>,
// Just for testing
id: usize,
}
impl ArcVec {
pub fn new(data: Arc<Vec<u8>>) -> ArcVec {
ArcVec {
data,
id: COUNTER.fetch_add(1, Ordering::SeqCst),
}
}
}
#[pyproto]
impl pyo3::class::buffer::PyBufferProtocol for ArcVec {
fn bf_getbuffer(
slf: PyRefMut<Self>,
view: *mut ffi::Py_buffer,
flags: std::os::raw::c_int,
) -> PyResult<()> {
println!("getbuffer {}", slf.id);
if view.is_null() {
return Err(pyo3::exceptions::PyBufferError::new_err("View is null"));
}
if (flags & ffi::PyBUF_WRITABLE) == ffi::PyBUF_WRITABLE {
return Err(pyo3::exceptions::PyBufferError::new_err(
"Object is not writable",
));
}
unsafe {
(*view).obj = slf.as_ptr();
ffi::Py_INCREF((*view).obj);
}
let bytes = slf.data.as_ref();
unsafe {
(*view).buf = bytes.as_ptr() as *mut std::os::raw::c_void;
(*view).len = bytes.len() as isize;
(*view).readonly = 1;
(*view).itemsize = 1;
(*view).format = std::ptr::null_mut();
if (flags & ffi::PyBUF_FORMAT) == ffi::PyBUF_FORMAT {
let msg = std::ffi::CStr::from_bytes_with_nul(b"B\0").unwrap();
(*view).format = msg.as_ptr() as *mut _;
}
(*view).ndim = 1;
(*view).shape = std::ptr::null_mut();
if (flags & ffi::PyBUF_ND) == ffi::PyBUF_ND {
(*view).shape = (&((*view).len)) as *const _ as *mut _;
}
(*view).strides = std::ptr::null_mut();
if (flags & ffi::PyBUF_STRIDES) == ffi::PyBUF_STRIDES {
(*view).strides = &((*view).itemsize) as *const _ as *mut _;
}
(*view).suboffsets = std::ptr::null_mut();
(*view).internal = std::ptr::null_mut();
}
Ok(())
}
fn bf_releasebuffer(slf: PyRefMut<Self>, _view: *mut ffi::Py_buffer) {
println!("releasebuffer {}", slf.id);
}
}
impl Drop for ArcVec {
fn drop(&mut self) {
println!("Dropped {}", &self.id);
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment