fix and test

Author: flying-sheep

src/npyffi/random.rs

@@ -6,13 +6,13 @@ use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyCapsule};
 #[derive(Debug, Clone, Copy)]
 pub struct npy_bitgen {
     pub state: *mut c_void,
-    pub next_uint64: NonNull<unsafe extern "C" fn(*mut c_void) -> super::npy_uint64>, //nogil
-    pub next_uint32: NonNull<unsafe extern "C" fn(*mut c_void) -> super::npy_uint32>, //nogil
-    pub next_double: NonNull<unsafe extern "C" fn(*mut c_void) -> libc::c_double>, //nogil
-    pub next_raw: NonNull<unsafe extern "C" fn(*mut c_void) -> super::npy_uint64>, //nogil
+    pub next_uint64: unsafe extern "C" fn(*mut c_void) -> super::npy_uint64, //nogil
+    pub next_uint32: unsafe extern "C" fn(*mut c_void) -> super::npy_uint32, //nogil
+    pub next_double: unsafe extern "C" fn(*mut c_void) -> libc::c_double, //nogil
+    pub next_raw: unsafe extern "C" fn(*mut c_void) -> super::npy_uint64, //nogil
 }
 
-pub fn get_bitgen_api<'py>(bitgen: Bound<'py, PyAny>) -> PyResult<NonNull<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"));
     let ptr = capsule.pointer() as *mut npy_bitgen;

src/random.rs

@@ -1,14 +1,12 @@
-//! Safe interface for NumPy's random [`BitGenerator`][]
-//!
-//! `BitGenerator`: https://numpy.org/doc/stable//reference/random/bit_generators/generated/numpy.random.BitGenerator.html
+//! Safe interface for NumPy's random [`BitGenerator`]
 
 use pyo3::{ffi, prelude::*, sync::GILOnceCell, types::PyType, PyTypeInfo};
 
 use crate::npyffi::get_bitgen_api;
 
-///! Wrapper for NumPy's random [`BitGenerator`][]
-/// 
-///! [BitGenerator]: https://numpy.org/doc/stable/reference/random/bit_generators/generated/numpy.random.BitGenerator.html
+///! Wrapper for NumPy's random [`BitGenerator`][bg]
+///!
+///! [bg]: https://numpy.org/doc/stable//reference/random/bit_generators/generated/numpy.random.BitGenerator.html
 #[repr(transparent)]
 pub struct BitGenerator(PyAny);
 
@@ -33,3 +31,50 @@ unsafe impl PyTypeInfo for BitGenerator {
     }
 }
 
+/// Methods for [`BitGenerator`]
+pub trait BitGeneratorMethods {
+    /// Returns the next random unsigned 64 bit integer
+    fn next_uint64(&self) -> u64;
+    /// Returns the next random unsigned 32 bit integer
+    fn next_uint32(&self) -> u32;
+    /// Returns the next random double
+    fn next_double(&self) -> libc::c_double;
+    /// Returns the next raw value (can be used for testing)
+    fn next_raw(&self) -> u64;
+}
+
+// TODO: cache npy_bitgen pointer
+impl<'py> BitGeneratorMethods for Bound<'py, BitGenerator> {
+    fn next_uint64(&self) -> u64 {
+        todo!()
+    }
+    fn next_uint32(&self) -> u32 {
+        todo!()
+    }
+    fn next_double(&self) -> libc::c_double {
+        todo!()
+    }
+    fn next_raw(&self) -> u64 {
+        let mut api = get_bitgen_api(self.as_any()).expect("Could not get bitgen");
+        unsafe {
+            let api = api.as_mut();
+            (api.next_raw)(api.state)
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    
+    #[test]
+    fn test_bitgen() -> PyResult<()> {
+        Python::with_gil(|py| {
+            let default_rng = py.import("numpy.random")?.getattr("default_rng")?;
+            let bitgen = default_rng.call0()?.getattr("bit_generator")?.downcast_into::<BitGenerator>()?;
+            let res = bitgen.next_raw();
+            dbg!(res);
+            Ok(())
+        })
+    }   
+}