BitGenerator support

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "[rust]": {
+        "editor.defaultFormatter": "rust-lang.rust-analyzer",
+        "editor.formatOnSave": true,
+    },
+    "rust-analyzer.cargo.features": "all",
+}

Cargo.toml

@@ -23,6 +23,7 @@ num-integer = "0.1"
 num-traits = "0.2"
 ndarray = ">= 0.15, < 0.17"
 pyo3 = { version = "0.25.0", default-features = false, features = ["macros"] }
+rand = { version = "0.9.1", default-features = false, optional = true }
 rustc-hash = "2.0"
 
 [dev-dependencies]

examples/simple/src/lib.rs

@@ -113,7 +113,7 @@ fn rust_ext<'py>(m: &Bound<'py, PyModule>) -> PyResult<()> {
     // This crate follows a strongly-typed approach to wrapping NumPy arrays
     // while Python API are often expected to work with multiple element types.
     //
-    // That kind of limited polymorphis can be recovered by accepting an enumerated type
+    // That kind of limited polymorphism can be recovered by accepting an enumerated type
     // covering the supported element types and dispatching into a generic implementation.
     #[derive(FromPyObject)]
     enum SupportedArray<'py> {

src/lib.rs

@@ -79,6 +79,7 @@ pub mod datetime;
 mod dtype;
 mod error;
 pub mod npyffi;
+pub mod random;
 mod slice_container;
 mod strings;
 mod sum_products;

src/npyffi/mod.rs

@@ -94,11 +94,13 @@ macro_rules! impl_api {
 pub mod array;
 pub mod flags;
 pub mod objects;
+pub mod random;
 pub mod types;
 pub mod ufunc;
 
 pub use self::array::*;
 pub use self::flags::*;
 pub use self::objects::*;
+pub use self::random::*;
 pub use self::types::*;
 pub use self::ufunc::*;

src/npyffi/objects.rs

@@ -1,4 +1,4 @@
-//! Low-Lebel binding for NumPy C API C-objects
+//! Low-Level binding for NumPy C API C-objects
 //!
 //! <https://numpy.org/doc/stable/reference/c-api/types-and-structures.html>
 #![allow(non_camel_case_types)]

src/npyffi/random.rs

@@ -0,0 +1,11 @@
+use std::ffi::c_void;
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy)] // TODO: can it be Clone and/or Copy?
+pub struct npy_bitgen {
+    pub state: *mut c_void,
+    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
+}

src/random.rs

@@ -0,0 +1,238 @@
+//! Safe interface for NumPy's random [`BitGenerator`][bg].
+//!
+//! Using the patterns described in [“Extending `numpy.random`”][ext],
+//! you can generate random numbers without holding the GIL,
+//! by [acquiring][`PyBitGeneratorMethods::lock`] a lock [guard][`PyBitGeneratorGuard`] for the [`PyBitGenerator`]:
+//!
+//! ```rust
+//! use pyo3::prelude::*;
+//! use numpy::random::{PyBitGenerator, PyBitGeneratorMethods as _};
+//!
+//! let mut bitgen = Python::with_gil(|py| -> PyResult<_> {
+//!     let default_rng = py.import("numpy.random")?.getattr("default_rng")?.call0()?;
+//!     let bit_generator = default_rng.getattr("bit_generator")?.downcast_into::<PyBitGenerator>()?;
+//!     bit_generator.lock()
+//! })?;
+//! let random_number = bitgen.next_u64();
+//! ```
+//!
+//! With the [`rand`] crate installed, you can also use the [`rand::Rng`] APIs from the [`PyBitGeneratorGuard`]:
+//!
+//! ```rust
+//! use rand::Rng as _;
+//!
+//! if bitgen.random_ratio(1, 1_000_000) {
+//!     println!("a sure thing");
+//! }
+//! ```
+//!
+//! [bg]: https://numpy.org/doc/stable//reference/random/bit_generators/generated/numpy.random.BitGenerator.html
+//! [ext]: https://numpy.org/doc/stable/reference/random/extending.html
+
+use std::ptr::NonNull;
+
+use pyo3::{
+    exceptions::PyRuntimeError,
+    ffi,
+    prelude::*,
+    sync::GILOnceCell,
+    types::{DerefToPyAny, PyCapsule, PyType},
+    PyTypeInfo,
+};
+
+use crate::npyffi::npy_bitgen;
+
+/// Wrapper for [`np.random.BitGenerator`][bg].
+///
+/// See also [`PyBitGeneratorMethods`].
+///
+/// [bg]: https://numpy.org/doc/stable//reference/random/bit_generators/generated/numpy.random.BitGenerator.html
+#[repr(transparent)]
+pub struct PyBitGenerator(PyAny);
+
+impl DerefToPyAny for PyBitGenerator {}
+
+unsafe impl PyTypeInfo for PyBitGenerator {
+    const NAME: &'static str = "PyBitGenerator";
+    const MODULE: Option<&'static str> = Some("numpy.random");
+
+    fn type_object_raw<'py>(py: Python<'py>) -> *mut ffi::PyTypeObject {
+        const CLS: GILOnceCell<Py<PyType>> = GILOnceCell::new();
+        let cls = CLS
+            .get_or_try_init::<_, PyErr>(py, || {
+                Ok(py
+                    .import("numpy.random")?
+                    .getattr("BitGenerator")?
+                    .downcast_into::<PyType>()?
+                    .unbind())
+            })
+            .expect("Failed to get BitGenerator type object")
+            .clone_ref(py)
+            .into_bound(py);
+        cls.as_type_ptr()
+    }
+}
+
+/// Methods for [`PyBitGenerator`].
+pub trait PyBitGeneratorMethods {
+    /// Acquire a lock on the BitGenerator to allow calling its methods in.
+    fn lock(&self) -> PyResult<PyBitGeneratorGuard>;
+}
+
+impl<'py> PyBitGeneratorMethods for Bound<'py, PyBitGenerator> {
+    fn lock(&self) -> PyResult<PyBitGeneratorGuard> {
+        let capsule = self.getattr("capsule")?.downcast_into::<PyCapsule>()?;
+        let lock = self.getattr("lock")?;
+        if lock.getattr("locked")?.call0()?.extract()? {
+            return Err(PyRuntimeError::new_err("BitGenerator is already locked"));
+        }
+        lock.getattr("acquire")?.call0()?;
+
+        assert_eq!(capsule.name()?, Some(c"BitGenerator"));
+        let ptr = capsule.pointer() as *mut npy_bitgen;
+        let non_null = match NonNull::new(ptr) {
+            Some(non_null) => non_null,
+            None => {
+                lock.getattr("release")?.call0()?;
+                return Err(PyRuntimeError::new_err("Invalid BitGenerator capsule"));
+            }
+        };
+        Ok(PyBitGeneratorGuard {
+            raw_bitgen: non_null,
+            lock: lock.unbind(),
+        })
+    }
+}
+
+impl<'py> TryFrom<&Bound<'py, PyBitGenerator>> for PyBitGeneratorGuard {
+    type Error = PyErr;
+    fn try_from(value: &Bound<'py, PyBitGenerator>) -> Result<Self, Self::Error> {
+        value.lock()
+    }
+}
+
+/// [`PyBitGenerator`] lock allowing to access its methods without holding the GIL.
+pub struct PyBitGeneratorGuard {
+    raw_bitgen: NonNull<npy_bitgen>,
+    lock: Py<PyAny>,
+}
+
+// SAFETY: We hold the `BitGenerator.lock`,
+// so nothing apart from us is allowed to change its state.
+impl PyBitGeneratorGuard {
+    /// Returns the next random unsigned 64 bit integer.
+    pub fn next_uint64(&mut self) -> u64 {
+        unsafe {
+            let bitgen = self.raw_bitgen.as_mut();
+            (bitgen.next_uint64)(bitgen.state)
+        }
+    }
+    /// Returns the next random unsigned 32 bit integer.
+    pub fn next_uint32(&mut self) -> u32 {
+        unsafe {
+            let bitgen = self.raw_bitgen.as_mut();
+            (bitgen.next_uint32)(bitgen.state)
+        }
+    }
+    /// Returns the next random double.
+    pub fn next_double(&mut self) -> libc::c_double {
+        unsafe {
+            let bitgen = self.raw_bitgen.as_mut();
+            (bitgen.next_double)(bitgen.state)
+        }
+    }
+    /// Returns the next raw value (can be used for testing).
+    pub fn next_raw(&mut self) -> u64 {
+        unsafe {
+            let bitgen = self.raw_bitgen.as_mut();
+            (bitgen.next_raw)(bitgen.state)
+        }
+    }
+}
+
+impl Drop for PyBitGeneratorGuard {
+    fn drop(&mut self) {
+        let r = Python::with_gil(|py| -> PyResult<()> {
+            self.lock.bind(py).getattr("release")?.call0()?;
+            Ok(())
+        });
+        if let Err(e) = r {
+            eprintln!("Failed to release BitGenerator lock: {e}");
+        }
+    }
+}
+
+#[cfg(feature = "rand")]
+impl rand::RngCore for PyBitGeneratorGuard {
+    fn next_u32(&mut self) -> u32 {
+        self.next_uint32()
+    }
+    fn next_u64(&mut self) -> u64 {
+        self.next_uint64()
+    }
+    fn fill_bytes(&mut self, dst: &mut [u8]) {
+        rand::rand_core::impls::fill_bytes_via_next(self, dst)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn get_bit_generator<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyBitGenerator>> {
+        let default_rng = py.import("numpy.random")?.getattr("default_rng")?.call0()?;
+        let bit_generator = default_rng
+            .getattr("bit_generator")?
+            .downcast_into::<PyBitGenerator>()?;
+        Ok(bit_generator)
+    }
+
+    #[test]
+    fn bitgen() -> PyResult<()> {
+        let mut bitgen = Python::with_gil(|py| get_bit_generator(py)?.lock())?;
+        let _ = bitgen.next_raw();
+        std::mem::drop(bitgen);
+        Ok(())
+    }
+
+    /// Test that the `rand::Rng` APIs work
+    #[cfg(feature = "rand")]
+    #[test]
+    fn rand() -> PyResult<()> {
+        use rand::Rng as _;
+
+        let mut bitgen = Python::with_gil(|py| get_bit_generator(py)?.lock())?;
+        assert!(bitgen.random_ratio(1, 1));
+        assert!(!bitgen.random_ratio(0, 1));
+        std::mem::drop(bitgen);
+        Ok(())
+    }
+
+    /// Test that releasing the lock works while holding the GIL
+    #[test]
+    fn unlock_with_held_gil() -> PyResult<()> {
+        Python::with_gil(|py| {
+            let generator = get_bit_generator(py)?;
+            let mut bitgen = generator.lock()?;
+            let _ = bitgen.next_raw();
+            std::mem::drop(bitgen);
+            assert!(!generator
+                .getattr("lock")?
+                .getattr("locked")?
+                .call0()?
+                .extract()?);
+            Ok(())
+        })
+    }
+
+    #[test]
+    fn double_lock_fails() -> PyResult<()> {
+        Python::with_gil(|py| {
+            let generator = get_bit_generator(py)?;
+            let d1 = generator.lock()?;
+            assert!(generator.lock().is_err());
+            std::mem::drop(d1);
+            Ok(())
+        })
+    }
+}