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,264 @@
+//! 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`]:
+//!
+//! ```
+//! use pyo3::prelude::*;
+//! use numpy::random::{PyBitGenerator, PyBitGeneratorMethods as _};
+//!
+//! fn default_bit_gen<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyBitGenerator>> {
+//!     let default_rng = py.import("numpy.random")?.call_method0("default_rng")?;
+//!     let bit_generator = default_rng.getattr("bit_generator")?.downcast_into()?;
+//!     Ok(bit_generator)
+//! }
+//!
+//! let random_number = Python::with_gil(|py| -> PyResult<_> {
+//!     let mut bitgen = default_bit_gen(py)?.lock()?;
+//!     // use bitgen without holding the GIL
+//!     Ok(py.allow_threads(|| bitgen.next_uint64()))
+//! })?;
+//! # Ok::<(), PyErr>(())
+//! ```
+//!
+//! With the [`rand`] crate installed, you can also use the [`rand::Rng`] APIs from the [`PyBitGeneratorGuard`]:
+//!
+//! ```
+//! # use pyo3::prelude::*;
+//! use rand::Rng as _;
+//! # use numpy::random::{PyBitGenerator, PyBitGeneratorMethods as _};
+//! # // TODO: reuse function definition from above?
+//! # fn default_bit_gen<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyBitGenerator>> {
+//! #     let default_rng = py.import("numpy.random")?.call_method0("default_rng")?;
+//! #     let bit_generator = default_rng.getattr("bit_generator")?.downcast_into()?;
+//! #     Ok(bit_generator)
+//! # }
+//!
+//! Python::with_gil(|py| -> PyResult<_> {
+//!     let mut bitgen = default_bit_gen(py)?.lock()?;
+//!     if bitgen.random_ratio(1, 1_000_000) {
+//!         println!("a sure thing");
+//!     }
+//!     Ok(())
+//! })?;
+//! # Ok::<(), PyErr>(())
+//! ```
+//!
+//! [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<'py> {
+    /// Acquire a lock on the BitGenerator to allow calling its methods in.
+    fn lock(&self) -> PyResult<PyBitGeneratorGuard<'py>>;
+}
+
+impl<'py> PyBitGeneratorMethods<'py> for Bound<'py, PyBitGenerator> {
+    fn lock(&self) -> PyResult<PyBitGeneratorGuard<'py>> {
+        let capsule = self.getattr("capsule")?.downcast_into::<PyCapsule>()?;
+        let lock = self.getattr("lock")?;
+        if lock.call_method0("locked")?.extract()? {
+            return Err(PyRuntimeError::new_err("BitGenerator is already locked"));
+        }
+        lock.call_method0("acquire")?;
+
+        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.call_method0("release")?;
+                return Err(PyRuntimeError::new_err("Invalid BitGenerator capsule"));
+            }
+        };
+        Ok(PyBitGeneratorGuard {
+            raw_bitgen: non_null,
+            _capsule: capsule.unbind(),
+            lock: lock.unbind(),
+            py: self.py(),
+        })
+    }
+}
+
+impl<'py> TryFrom<&Bound<'py, PyBitGenerator>> for PyBitGeneratorGuard<'py> {
+    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<'py> {
+    raw_bitgen: NonNull<npy_bitgen>,
+    /// This field makes sure the `raw_bitgen` inside the capsule doesn’t get deallocated.
+    _capsule: Py<PyCapsule>,
+    /// This lock makes sure no other threads try to use the BitGenerator while we do.
+    lock: Py<PyAny>,
+    /// This should be an unsafe field (https://github.com/rust-lang/rust/issues/132922)
+    ///
+    /// SAFETY: only use this in `Drop::drop` (when we are sure the GIL is held).
+    py: Python<'py>,
+}
+
+// SAFETY: we can’t have public APIs that access the Python objects,
+// only the `raw_bitgen` pointer.
+unsafe impl Send for PyBitGeneratorGuard<'_> {}
+
+impl Drop for PyBitGeneratorGuard<'_> {
+    fn drop(&mut self) {
+        // ignore errors. This includes when `try_drop` was called manually
+        let _ = self.lock.bind(self.py).call_method0("release");
+    }
+}
+
+// SAFETY: We hold the `BitGenerator.lock`,
+// so nothing apart from us is allowed to change its state.
+impl<'py> PyBitGeneratorGuard<'py> {
+    /// Drop the lock manually before `Drop::drop` tries to do it (used for testing).
+    #[allow(dead_code)]
+    fn try_drop(self, py: Python<'py>) -> PyResult<()> {
+        self.lock.bind(py).call_method0("release")?;
+        Ok(())
+    }
+
+    /// Returns the next random unsigned 64 bit integer.
+    pub fn next_uint64(&mut self) -> u64 {
+        unsafe {
+            let bitgen = *self.raw_bitgen.as_ptr();
+            (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_ptr();
+            (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_ptr();
+            (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_ptr();
+            (bitgen.next_raw)(bitgen.state)
+        }
+    }
+}
+
+#[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")?.call_method0("default_rng")?;
+        let bit_generator = default_rng
+            .getattr("bit_generator")?
+            .downcast_into::<PyBitGenerator>()?;
+        Ok(bit_generator)
+    }
+
+    /// Test the primary use case: acquire the lock, release the GIL, then use the lock
+    #[test]
+    fn use_outside_gil() -> PyResult<()> {
+        Python::with_gil(|py| {
+            let mut bitgen = get_bit_generator(py)?.lock()?;
+            py.allow_threads(|| {
+                let _ = bitgen.next_raw();
+            });
+            assert!(bitgen.try_drop(py).is_ok());
+            Ok(())
+        })
+    }
+
+    /// Test that the `rand::Rng` APIs work
+    #[cfg(feature = "rand")]
+    #[test]
+    fn rand() -> PyResult<()> {
+        use rand::Rng as _;
+
+        Python::with_gil(|py| {
+            let mut bitgen = get_bit_generator(py)?.lock()?;
+            py.allow_threads(|| {
+                assert!(bitgen.random_ratio(1, 1));
+                assert!(!bitgen.random_ratio(0, 1));
+            });
+            assert!(bitgen.try_drop(py).is_ok());
+            Ok(())
+        })
+    }
+
+    #[test]
+    fn double_lock_fails() -> PyResult<()> {
+        Python::with_gil(|py| {
+            let generator = get_bit_generator(py)?;
+            let bitgen = generator.lock()?;
+            assert!(generator.lock().is_err());
+            assert!(bitgen.try_drop(py).is_ok());
+            Ok(())
+        })
+    }
+}