WIP bitgen

Author: flying-sheep

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "rust-analyzer.cargo.features": "all"
+}

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,19 @@
+use std::{ffi::c_void, ptr::NonNull};
+
+use pyo3::{prelude::*, types::PyCapsule};
+
+#[repr(C)]
+#[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 fn get_bitgen_api<'py>(bitgen: Bound<'py, PyAny>) -> PyResult<*mut npy_bitgen> {
+    let capsule = bitgen.getattr("capsule")?.downcast_into::<PyCapsule>()?;
+    assert_eq!(capsule.name()?, Some(c"BitGenerator"));
+    Ok(capsule.pointer() as *mut npy_bitgen)
+}

src/random.rs

@@ -0,0 +1,35 @@
+//! Safe interface for NumPy's random [`BitGenerator`][]
+//!
+//! `BitGenerator`: https://numpy.org/doc/stable//reference/random/bit_generators/generated/numpy.random.BitGenerator.html
+
+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
+#[repr(transparent)]
+pub struct BitGenerator(PyAny);
+
+unsafe impl PyTypeInfo for BitGenerator {
+    const NAME: &'static str = "BitGenerator";
+    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()
+    }
+}
+