From 33e9b9eb0fb1014f1c12176ec9e8fcb077ab5267 Mon Sep 17 00:00:00 2001
From: Benjamin Bolte <ben@bolte.cc>
Date: Fri, 20 Sep 2024 18:13:45 +0800
Subject: [PATCH] wip

---
 actuator/rust/Cargo.toml          |  5 +++--
 actuator/rust/src/bin/stub_gen.rs |  7 +++++++
 actuator/rust/src/lib.rs          | 16 ++++++++++++--
 setup.py                          | 35 ++++++++++++++++++-------------
 4 files changed, 44 insertions(+), 19 deletions(-)
 create mode 100644 actuator/rust/src/bin/stub_gen.rs

diff --git a/actuator/rust/Cargo.toml b/actuator/rust/Cargo.toml
index 9e843a8..e47b0a0 100644
--- a/actuator/rust/Cargo.toml
+++ b/actuator/rust/Cargo.toml
@@ -5,7 +5,8 @@ edition = "2021"
 
 [lib]
 name = "lib"
-crate-type = ["cdylib"]
+crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-pyo3 = { version = "0.17", features = ["extension-module"] }
+pyo3 = { version = "0.21", features = ["extension-module"] }
+pyo3-stub-gen = "0.6.0"
diff --git a/actuator/rust/src/bin/stub_gen.rs b/actuator/rust/src/bin/stub_gen.rs
new file mode 100644
index 0000000..e48e995
--- /dev/null
+++ b/actuator/rust/src/bin/stub_gen.rs
@@ -0,0 +1,7 @@
+use pyo3_stub_gen::Result;
+
+fn main() -> Result<()> {
+    let stub = pure::stub_info()?;
+    stub.generate()?;
+    Ok(())
+}
diff --git a/actuator/rust/src/lib.rs b/actuator/rust/src/lib.rs
index 9ec676b..b1b579d 100644
--- a/actuator/rust/src/lib.rs
+++ b/actuator/rust/src/lib.rs
@@ -1,7 +1,19 @@
 use pyo3::prelude::*;
+use pyo3_stub_gen::{define_stub_info_gatherer, derive::gen_stub_pyfunction};
 
-#[pymodule]
-fn lib(_py: Python, _m: &PyModule) -> PyResult<()> {
+#[pyfunction]
+fn hello_world() {
     println!("Hello, world!");
+}
+
+#[pyfunction]
+fn add(a: i64, b: i64) -> i64 {
+    a + b
+}
+
+#[pymodule]
+fn lib(py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(hello_world, m)?)?;
+    m.add_function(wrap_pyfunction!(add, m)?)?;
     Ok(())
 }
diff --git a/setup.py b/setup.py
index 25e1f44..8736746 100644
--- a/setup.py
+++ b/setup.py
@@ -2,25 +2,12 @@
 #!/usr/bin/env python
 """Setup script for the project."""
 
-import os
 import re
 import subprocess
 
-from setuptools import setup
-from setuptools.command.build_ext import build_ext as _build_ext
+from setuptools import Command, setup
 from setuptools_rust import Binding, RustExtension
 
-
-class BuildExtWithStubgen(_build_ext):
-    def run(self) -> None:
-        super().run()
-
-        module_name = "actuator.rust.lib"
-        output_dir = os.path.dirname(__file__)
-        os.makedirs(output_dir, exist_ok=True)
-        subprocess.run(["stubgen", "-m", module_name, "-o", output_dir], check=False)
-
-
 with open("README.md", "r", encoding="utf-8") as f:
     long_description: str = f.read()
 
@@ -39,6 +26,22 @@ def run(self) -> None:
 version: str = version_re.group(1)
 
 
+class PostInstallCommand(Command):
+    """Post-installation for installation mode."""
+
+    description = "Run stub_gen after installation"
+    user_options = []
+
+    def initialize_options(self) -> None:
+        pass
+
+    def finalize_options(self) -> None:
+        pass
+
+    def run(self) -> None:
+        subprocess.check_call(["cargo", "run", "--bin", "stub_gen"], cwd="actuator/rust")
+
+
 setup(
     name="actuator",
     version=version,
@@ -52,7 +55,6 @@ def run(self) -> None:
             binding=Binding.PyO3,
         )
     ],
-    cmdclass={"build_ext": BuildExtWithStubgen},
     setup_requires=[
         "setuptools-rust",
         "mypy",  # For stubgen
@@ -65,4 +67,7 @@ def run(self) -> None:
     install_requires=requirements,
     tests_require=requirements_dev,
     extras_require={"dev": requirements_dev},
+    cmdclass={
+        "post_install": PostInstallCommand,
+    },
 )