1 files changed, 90 insertions, 36 deletions

diff --git a/hittekaart-py/src/lib.rs b/hittekaart-py/src/lib.rs
index fe97aa7..2f99793 100644
--- a/hittekaart-py/src/lib.rs
+++ b/hittekaart-py/src/lib.rs
@@ -123,7 +123,7 @@ impl Storage {
 }
 
 impl Storage {
-    fn open(&self) -> PyResult<Box<dyn hittekaart::storage::Storage>> {
+    fn open(&self) -> PyResult<Box<dyn hittekaart::storage::Storage + Send>> {
         match self.0 {
             StorageType::Folder(ref path) => {
                 let storage = hittekaart::storage::Folder::new(path.clone());
@@ -194,13 +194,72 @@ impl TilehuntRenderer {
     }
 }
 
+/// Tile generation settings.
+///
+/// This contains everything that is overarching to the renderers and output modules, such as zoom
+/// levels and thread count.
+#[pyclass]
+struct Settings {
+    #[pyo3(get, set)]
+    /// Smallest zoom level that will be generated.
+    min_zoom: u32,
+
+    #[pyo3(get, set)]
+    /// Largest zoom level that will be generated.
+    max_zoom: u32,
+
+    #[pyo3(get, set)]
+    /// How many threads to use for generation.
+    ///
+    /// A count of 0 will automatically use as many threads as you have CPU cores.
+    threads: u32,
+}
+
+#[pymethods]
+impl Settings {
+    #[new]
+    #[pyo3(signature = (min_zoom = 1, max_zoom = 19, threads = 0))]
+    fn new(min_zoom: u32, max_zoom: u32, threads: u32) -> Settings {
+        Settings {
+            min_zoom,
+            max_zoom,
+            threads,
+        }
+    }
+
+    fn __repr__(&self) -> String {
+        format!(
+            "Settings(min_zoom={}, max_zoom={}, threads={})",
+            self.min_zoom, self.max_zoom, self.threads
+        )
+    }
+}
+
+macro_rules! dispatch_generate {
+    ($settings:expr, $tracks:expr, $renderer:expr, $storage:expr => <$type:ty>, $(<$types:ty>,)*) => {
+        if let Ok(r) = $renderer.downcast::<$type>() {
+            do_generate($settings, $tracks, &r.borrow().inner, $storage)
+        } else {
+            dispatch_generate!($settings, $tracks, $renderer, $storage => $(<$types>,)*)
+        }
+    };
+
+    ($settings:expr, $tracks:expr, $renderer:expr, $storage:expr =>) => {
+        Err(PyTypeError::new_err(
+            "Expected a HeatmapRenderer, MarktileRenderer or TilehuntRenderer",
+        ))
+    };
+}
+
 /// Generate the heatmap.
 ///
+/// * settings are the rendering Settings
 /// * items is an iterable of Track
 /// * renderer should be one of the renderers (such as HeatmapRenderer)
 /// * storage is the Storage output
 #[pyfunction]
 fn generate(
+    settings: &Bound<'_, Settings>,
     items: &Bound<'_, PyAny>,
     renderer: &Bound<'_, PyAny>,
     storage: &Bound<'_, Storage>,
@@ -212,53 +271,48 @@ fn generate(
         tracks.push(item.extract::<Track>()?.inner);
     }
 
-    if let Ok(r) = renderer.downcast::<HeatmapRenderer>() {
-        do_generate(tracks, &r.borrow().inner, &mut *storage.borrow().open()?)
-    } else if let Ok(r) = renderer.downcast::<MarktileRenderer>() {
-        do_generate(tracks, &r.borrow().inner, &mut *storage.borrow().open()?)
-    } else if let Ok(r) = renderer.downcast::<TilehuntRenderer>() {
-        do_generate(tracks, &r.borrow().inner, &mut *storage.borrow().open()?)
-    } else {
-        Err(PyTypeError::new_err("Expected a HeatmapRenderer, MarktileRenderer or TilehuntRenderer"))
+    let settings = &*settings.borrow();
+
+    // We cannot easily do dynamic dispatch here, because Renderer::Prepared exists. Maybe this can
+    // change in the future, but for now we have to stick to this:
+    dispatch_generate! {
+        settings, tracks, renderer, &mut *storage.borrow().open()? =>
+            <HeatmapRenderer>,
+            <MarktileRenderer>,
+            <TilehuntRenderer>,
     }
 }
 
 fn do_generate<R: Renderer>(
+    settings: &Settings,
     tracks: Vec<Vec<Coordinates>>,
     renderer: &R,
-    storage: &mut dyn hittekaart::storage::Storage,
+    storage: &mut (dyn hittekaart::storage::Storage + Send),
 ) -> PyResult<()> {
     storage.prepare().map_err(|e| err_to_py(&e))?;
 
-    for zoom in 0..=19 {
-        let counter =
-            renderer::prepare(renderer, zoom, &tracks, || Ok(())).map_err(|e| err_to_py(&e))?;
+    let pool = rayon::ThreadPoolBuilder::new()
+        .num_threads(settings.threads.try_into().unwrap())
+        .build()
+        .map_err(|e| err_to_py(&e))?;
 
-        storage.prepare_zoom(zoom).map_err(|e| err_to_py(&e))?;
+    pool.install(|| {
+        for zoom in settings.min_zoom..=settings.max_zoom {
+            let counter =
+                renderer::prepare(renderer, zoom, &tracks, || Ok(())).map_err(|e| err_to_py(&e))?;
 
-        renderer::colorize(renderer, counter, |rendered_tile| {
-            storage.store(zoom, rendered_tile.x, rendered_tile.y, &rendered_tile.data)?;
-            Ok(())
-        })
-        .map_err(|e| err_to_py(&e))?;
-    }
-    storage.finish().map_err(|e| err_to_py(&e))?;
+            storage.prepare_zoom(zoom).map_err(|e| err_to_py(&e))?;
 
-    Ok(())
-}
+            renderer::colorize(renderer, counter, |rendered_tile| {
+                storage.store(zoom, rendered_tile.x, rendered_tile.y, &rendered_tile.data)?;
+                Ok(())
+            })
+            .map_err(|e| err_to_py(&e))?;
+        }
+        storage.finish().map_err(|e| err_to_py(&e))?;
 
-/// Set the number of threads that hittekaart will use.
-///
-/// Note that this is a global function, it will affect all subsequent calls.
-///
-/// Note further that you may only call this function once, at startup. Calls after the thread pool
-/// has been initialized (e.g. via a generate() or set_threads() call) will raise an exception.
-#[pyfunction]
-fn set_threads(threads: usize) -> PyResult<()> {
-    rayon::ThreadPoolBuilder::new()
-        .num_threads(threads)
-        .build_global()
-        .map_err(|e| err_to_py(&e))
+        Ok(())
+    })
 }
 
 /// Python bindings for the hittekaart heatmap tile generator.
@@ -278,8 +332,8 @@ fn hittekaart_py(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MarktileRenderer>()?;
     m.add_class::<TilehuntRenderer>()?;
     m.add_class::<Storage>()?;
+    m.add_class::<Settings>()?;
     m.add_function(wrap_pyfunction!(generate, m)?)?;
-    m.add_function(wrap_pyfunction!(set_threads, m)?)?;
     m.add("HitteError", py.get_type::<HitteError>())?;
     Ok(())
 }