2 files changed, 10 insertions, 0 deletions
diff --git a/hittekaart-py/Cargo.toml b/hittekaart-py/Cargo.toml
index 9a44118..7e69ef8 100644
--- a/hittekaart-py/Cargo.toml
+++ b/hittekaart-py/Cargo.toml
@@ -11,3 +11,4 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3 = "0.24.0"
 hittekaart = { path = "../hittekaart" }
+rayon = "1.10.0"
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(())
 }