expose function to set thread count

1 files changed, 9 insertions, 0 deletions

diff --git a/hittekaart-py/src/lib.rs b/hittekaart-py/src/lib.rs
index 4b67777..2ad0e56 100644
--- a/hittekaart-py/src/lib.rs
+++ b/hittekaart-py/src/lib.rs
@@ -164,6 +164,14 @@ fn do_generate<R: Renderer>(
     Ok(())
 }
 
+#[pyfunction]
+fn set_threads(threads: usize) -> PyResult<()> {
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(threads)
+        .build_global()
+        .map_err(|e| err_to_py(&e))
+}
+
 /// A Python module implemented in Rust.
 #[pymodule]
 fn hittekaart_py(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
@@ -171,6 +179,7 @@ fn hittekaart_py(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<HeatmapRenderer>()?;
     m.add_class::<Storage>()?;
     m.add_function(wrap_pyfunction!(generate, m)?)?;
+    m.add_function(wrap_pyfunction!(set_threads, m)?)?;
     m.add("HitteError", py.get_type::<HitteError>())?;
     Ok(())
 }