nonnull

Author: flying-sheep

src/npyffi/random.rs

@@ -1,6 +1,6 @@
 use std::{ffi::c_void, ptr::NonNull};
 
-use pyo3::{prelude::*, types::PyCapsule};
+use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyCapsule};
 
 #[repr(C)]
 #[derive(Debug, Clone, Copy)]
@@ -12,8 +12,9 @@ pub struct npy_bitgen {
     pub next_raw: NonNull<unsafe extern "C" fn(*mut c_void) -> super::npy_uint64>, //nogil
 }
 
-pub fn get_bitgen_api<'py>(bitgen: Bound<'py, PyAny>) -> PyResult<*mut npy_bitgen> {
+pub fn get_bitgen_api<'py>(bitgen: Bound<'py, PyAny>) -> PyResult<NonNull<npy_bitgen>> {
     let capsule = bitgen.getattr("capsule")?.downcast_into::<PyCapsule>()?;
     assert_eq!(capsule.name()?, Some(c"BitGenerator"));
-    Ok(capsule.pointer() as *mut npy_bitgen)
+    let ptr = capsule.pointer() as *mut npy_bitgen;
+    NonNull::new(ptr).ok_or_else(|| PyRuntimeError::new_err("Invalid BitGenerator capsule"))
 }