proc-macro added for pyo3-stub-gen to generate stub pyi

Author: unkcpz

python/Cargo.toml

@@ -11,7 +11,7 @@ readme = "README.md"
 
 [lib]
 name = "egobox"
-crate-type = ["cdylib"]
+crate-type = ["cdylib", "rlib"]
 
 [features]
 default = []
@@ -48,3 +48,8 @@ serde_json.workspace = true
 ctrlc.workspace = true
 
 argmin_testfunctions.workspace = true
+pyo3-stub-gen = { path = "../../pyo3-stub-gen/pyo3-stub-gen", features = ["numpy"] }
+
+[[bin]]
+name = "stub_gen"
+doc = false

python/py.typed

@@ -0,0 +1,8 @@
+use pyo3_stub_gen::Result;
+
+fn main() -> Result<()> {
+    env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
+    let stub = egobox::stub_info()?;
+    stub.generate()?;
+    Ok(())
+}

python/src/bin/stub_gen.rs

@@ -16,6 +16,7 @@ use ndarray::{concatenate, Array1, Array2, ArrayView2, Axis};
 use numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray2, ToPyArray};
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyfunction, gen_stub_pymethods};
 
 /// Utility function converting `xlimits` float data list specifying bounds of x components
 /// to x specified as a list of XType.Float types [egobox.XType]
@@ -25,6 +26,7 @@ use pyo3::prelude::*;
 ///
 /// # Returns
 ///     xtypes: nx-size list of XSpec(XType(FLOAT), [lower_bound, upper_bounds]) where `nx` is the dimension of x
+#[gen_stub_pyfunction]
 #[pyfunction]
 pub(crate) fn to_specs(py: Python, xlimits: Vec<Vec<f64>>) -> PyResult<PyObject> {
     if xlimits.is_empty() || xlimits[0].is_empty() {
@@ -149,6 +151,7 @@ pub(crate) fn to_specs(py: Python, xlimits: Vec<Vec<f64>>) -> PyResult<PyObject>
 ///     seed (int >= 0)
 ///         Random generator seed to allow computation reproducibility.
 ///      
+#[gen_stub_pyclass]
 #[pyclass]
 pub(crate) struct Egor {
     pub xspecs: PyObject,
@@ -174,6 +177,7 @@ pub(crate) struct Egor {
     pub seed: Option<u64>,
 }
 
+#[gen_stub_pyclass]
 #[pyclass]
 pub(crate) struct OptimResult {
     #[pyo3(get)]
@@ -186,6 +190,7 @@ pub(crate) struct OptimResult {
     y_doe: Py<PyArray2<f64>>,
 }
 
+#[gen_stub_pymethods]
 #[pymethods]
 impl Egor {
     #[new]

python/src/egor.rs

@@ -22,6 +22,7 @@ use ndarray::{Array1, Array2, Axis, Ix1, Ix2, Zip};
 use ndarray_rand::rand::SeedableRng;
 use numpy::{IntoPyArray, PyArray1, PyArray2, PyReadonlyArray2, PyReadonlyArrayDyn};
 use pyo3::prelude::*;
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 use rand_xoshiro::Xoshiro256Plus;
 
 /// Gaussian processes mixture builder
@@ -69,6 +70,7 @@ use rand_xoshiro::Xoshiro256Plus;
 ///     seed (int >= 0)
 ///         Random generator seed to allow computation reproducibility.
 ///         
+#[gen_stub_pyclass]
 #[pyclass]
 pub(crate) struct GpMix {
     pub n_clusters: usize,
@@ -82,6 +84,7 @@ pub(crate) struct GpMix {
     pub seed: Option<u64>,
 }
 
+#[gen_stub_pymethods]
 #[pymethods]
 impl GpMix {
     #[new]
@@ -218,9 +221,11 @@ impl GpMix {
 }
 
 /// A trained Gaussian processes mixture
+#[gen_stub_pyclass]
 #[pyclass]
 pub(crate) struct Gpx(Box<GpMixture>);
 
+#[gen_stub_pymethods]
 #[pymethods]
 impl Gpx {
     /// Get Gaussian processes mixture builder aka `GpMix`

python/src/gp_mix.rs

@@ -14,6 +14,7 @@ use types::*;
 
 use env_logger::{Builder, Env};
 use pyo3::prelude::*;
+use pyo3_stub_gen::define_stub_info_gatherer;
 
 #[doc(hidden)]
 #[pymodule]
@@ -55,3 +56,6 @@ fn egobox(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
 
     Ok(())
 }
+
+// Define a function to gather stub information.
+define_stub_info_gatherer!(stub_info);

python/src/lib.rs

@@ -3,7 +3,9 @@ use egobox_doe::{LhsKind, SamplingMethod};
 use egobox_ego::gpmix::mixint::MixintContext;
 use numpy::{IntoPyArray, PyArray2};
 use pyo3::prelude::*;
+use pyo3_stub_gen::derive::{gen_stub_pyclass_enum, gen_stub_pyfunction};
 
+#[gen_stub_pyclass_enum]
 #[pyclass(eq, eq_int, rename_all = "SCREAMING_SNAKE_CASE")]
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub enum Sampling {
@@ -27,6 +29,7 @@ pub enum Sampling {
 /// # Returns
 ///    ndarray of shape (n_samples, n_variables)
 ///
+#[gen_stub_pyfunction]
 #[pyfunction]
 #[pyo3(signature = (method, xspecs, n_samples, seed=None))]
 pub fn sampling(
@@ -89,6 +92,7 @@ pub fn sampling(
 /// # Returns
 ///    ndarray of shape (n_samples, n_variables)
 ///
+#[gen_stub_pyfunction]
 #[pyfunction]
 #[pyo3(signature = (xspecs, n_samples, seed=None))]
 pub(crate) fn lhs(

python/src/sampling.rs

@@ -21,6 +21,7 @@ use ndarray::{Array1, Array2, Axis, Ix1, Ix2, Zip};
 use ndarray_rand::rand::SeedableRng;
 use numpy::{IntoPyArray, PyArray1, PyArray2, PyReadonlyArray2};
 use pyo3::prelude::*;
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 use rand_xoshiro::Xoshiro256Plus;
 
 /// Sparse Gaussian processes mixture builder
@@ -58,6 +59,7 @@ use rand_xoshiro::Xoshiro256Plus;
 ///     seed (int >= 0)
 ///         Random generator seed to allow computation reproducibility.
 ///         
+#[gen_stub_pyclass]
 #[pyclass]
 pub(crate) struct SparseGpMix {
     pub correlation_spec: CorrelationSpec,
@@ -71,6 +73,7 @@ pub(crate) struct SparseGpMix {
     pub seed: Option<u64>,
 }
 
+#[gen_stub_pymethods]
 #[pymethods]
 impl SparseGpMix {
     #[new]
@@ -216,9 +219,11 @@ impl SparseGpMix {
 }
 
 /// A trained Gaussian processes mixture
+#[gen_stub_pyclass]
 #[pyclass]
 pub(crate) struct SparseGpx(Box<GpMixture>);
 
+#[gen_stub_pymethods]
 #[pymethods]
 impl SparseGpx {
     /// Get Gaussian processes mixture builder aka `GpSparse`

python/src/sparse_gp_mix.rs

@@ -1,5 +1,7 @@
 use pyo3::prelude::*;
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods};
 
+#[gen_stub_pyclass_enum]
 #[pyclass(eq, eq_int, rename_all = "UPPERCASE")]
 #[derive(Debug, Clone, PartialEq)]
 pub enum Recombination {
@@ -12,10 +14,12 @@ pub enum Recombination {
     Smooth = 1,
 }
 
+#[gen_stub_pyclass]
 #[pyclass]
 #[derive(Clone)]
 pub(crate) struct RegressionSpec(pub(crate) u8);
 
+#[gen_stub_pymethods]
 #[pymethods]
 impl RegressionSpec {
     #[classattr]
@@ -28,10 +32,12 @@ impl RegressionSpec {
     pub(crate) const QUADRATIC: u8 = egobox_moe::RegressionSpec::QUADRATIC.bits();
 }
 
+#[gen_stub_pyclass]
 #[pyclass]
 #[derive(Clone)]
 pub(crate) struct CorrelationSpec(pub(crate) u8);
 
+#[gen_stub_pymethods]
 #[pymethods]
 impl CorrelationSpec {
     #[classattr]
@@ -48,6 +54,7 @@ impl CorrelationSpec {
     pub(crate) const MATERN52: u8 = egobox_moe::CorrelationSpec::MATERN52.bits();
 }
 
+#[gen_stub_pyclass_enum]
 #[pyclass(eq, eq_int, rename_all = "UPPERCASE")]
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub(crate) enum InfillStrategy {
@@ -56,6 +63,7 @@ pub(crate) enum InfillStrategy {
     Wb2s = 3,
 }
 
+#[gen_stub_pyclass_enum]
 #[pyclass(eq, eq_int, rename_all = "UPPERCASE")]
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub(crate) enum ParInfillStrategy {
@@ -65,13 +73,15 @@ pub(crate) enum ParInfillStrategy {
     Clmin = 4,
 }
 
+#[gen_stub_pyclass_enum]
 #[pyclass(eq, eq_int, rename_all = "UPPERCASE")]
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub(crate) enum InfillOptimizer {
     Cobyla = 1,
     Slsqp = 2,
 }
 
+#[gen_stub_pyclass]
 #[pyclass]
 #[derive(Clone, Copy)]
 pub(crate) struct ExpectedOptimum {
@@ -93,6 +103,7 @@ impl ExpectedOptimum {
     }
 }
 
+#[gen_stub_pyclass_enum]
 #[pyclass(eq, eq_int, rename_all = "UPPERCASE")]
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub(crate) enum XType {
@@ -102,6 +113,7 @@ pub(crate) enum XType {
     Enum = 4,
 }
 
+#[gen_stub_pyclass]
 #[pyclass]
 #[derive(FromPyObject, Debug)]
 pub(crate) struct XSpec {
@@ -113,6 +125,7 @@ pub(crate) struct XSpec {
     pub(crate) tags: Vec<String>,
 }
 
+#[gen_stub_pymethods]
 #[pymethods]
 impl XSpec {
     #[new]
@@ -126,6 +139,7 @@ impl XSpec {
     }
 }
 
+#[gen_stub_pyclass_enum]
 #[pyclass(eq, eq_int, rename_all = "UPPERCASE")]
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub(crate) enum SparseMethod {