path: root/hittekaart-py/src/lib.rs

//! Python bindings for hittekaart.
//!
//! This module provides simple-to-use bindings for hittekaart generation from Python scripts.
use hittekaart::gpx::{self, Compression, Coordinates};
use hittekaart::renderer::{self, Renderer};
use pyo3::create_exception;
use pyo3::exceptions::PyTypeError;
use pyo3::prelude::*;
use std::error::Error;
use std::ffi::OsStr;
use std::fmt::Write as _;
use std::os::unix::ffi::OsStrExt as _;
use std::path::PathBuf;

create_exception!(hittekaart_py, HitteError, pyo3::exceptions::PyException);

/// Converts an error to a Python error.
///
/// This basically maps everything to [`HitteError`] and provides a stringified error explanation.
///
/// This recursively uses `::source()` to walk the chain.
fn err_to_py(mut error: &dyn Error) -> PyErr {
    let mut text = error.to_string();
    loop {
        match error.source() {
            None => break,
            Some(e) => error = e,
        }
        write!(&mut text, "\ncaused by: {error}").unwrap();
    }
    HitteError::new_err(text)
}

/// Represents a track.
///
/// This is what you want to load in order to render the heatmaps.
///
/// Tracks can be loaded from GPX files (see Track.from_file()) or from in-memory coordinates
/// (Track.from_coordinates()). Otherwise, tracks should be treated as opaque objects, whose only
/// purpose is to be passed to generate().
#[pyclass]
#[derive(Debug, Clone)]
struct Track {
    inner: Vec<Coordinates>,
}

#[pymethods]
impl Track {
    /// Load a track from a file.
    ///
    /// The path should be given as a bytes object.
    ///
    /// The compression parameter - if given - should be one of the strings "gzip" or "brotli". It
    /// can be set to None to use no compression.
    #[staticmethod]
    fn from_file(path: &[u8], compression: Option<&str>) -> PyResult<Track> {
        let compression = match compression {
            None | Some("") => Compression::None,
            Some("gzip") => Compression::Gzip,
            Some("brotli") => Compression::Brotli,
            Some(x) => return Err(HitteError::new_err(format!("invalid compression: {x}"))),
        };
        let track = gpx::extract_from_file(OsStr::from_bytes(path), compression)
            .map_err(|e| err_to_py(&e))?;
        Ok(Track { inner: track })
    }

    /// Load a track from the given coordinates.
    ///
    /// The coordinates should be a list of (longitude, latitude) tuples, where longitude and
    /// latitude are represented as floats.
    #[staticmethod]
    fn from_coordinates(coordinates: Vec<(f64, f64)>) -> Track {
        Track {
            inner: coordinates
                .iter()
                .map(|(lon, lat)| Coordinates::new(*lon, *lat))
                .collect(),
        }
    }

    fn __repr__(&self) -> String {
        format!("<Track [{} coordinates]>", self.inner.len())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum StorageType {
    Folder(PathBuf),
    Sqlite(PathBuf),
}

/// Represents a storage target.
///
/// Hittekaart can store tiles either in folders (easy-to-use, but wasteful), or as a SQLite
/// database (more space-efficient).
///
/// See hittekaart's README for more detailed information.
#[pyclass]
#[derive(Debug, Clone, PartialEq, Eq)]
struct Storage(StorageType);

#[pymethods]
impl Storage {
    /// Output to the given folder.
    ///
    /// This will create files path/{z}/{x}/{y}.png, where {z} is the zoom level, and {x} and {y}
    /// are the tile coordinates.
    #[staticmethod]
    #[pyo3(name = "Folder")]
    fn folder(path: &[u8]) -> Self {
        let path = OsStr::from_bytes(path);
        Storage(StorageType::Folder(path.into()))
    }

    /// Output to the given SQLite file.
    ///
    /// This will create a single table 'tiles' with the columns 'zoom', 'x', 'y' and 'data'.
    ///
    /// Note that you cannot "append" to an existing database, it must be a non-existing file.
    #[staticmethod]
    #[pyo3(name = "Sqlite")]
    fn sqlite(path: &[u8]) -> Self {
        let path = OsStr::from_bytes(path);
        Storage(StorageType::Sqlite(path.into()))
    }

    fn __repr__(&self) -> String {
        match self.0 {
            StorageType::Folder(ref path) => format!("<Storage.Folder path='{}'>", path.display()),
            StorageType::Sqlite(ref path) => format!("<Storage.Sqlite path='{}'>", path.display()),
        }
    }
}

impl 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());
                Ok(Box::new(storage))
            }
            StorageType::Sqlite(ref path) => {
                let storage = hittekaart::storage::Sqlite::connect(path.clone())
                    .map_err(|e| err_to_py(&e))?;
                Ok(Box::new(storage))
            }
        }
    }
}

/// A renderer that produces a heatmap.
///
/// The constructor takes no parameters: HeatmapRenderer()
#[pyclass]
struct HeatmapRenderer {
    inner: renderer::heatmap::Renderer,
}

#[pymethods]
impl HeatmapRenderer {
    #[new]
    fn new() -> HeatmapRenderer {
        HeatmapRenderer {
            inner: renderer::heatmap::Renderer,
        }
    }
}

/// A renderer that only marks visited tiles.
///
/// The constructor takes no parameters: MarktileRenderer()
#[pyclass]
struct MarktileRenderer {
    inner: renderer::marktile::Renderer,
}

#[pymethods]
impl MarktileRenderer {
    #[new]
    fn new() -> MarktileRenderer {
        MarktileRenderer {
            inner: renderer::marktile::Renderer,
        }
    }
}

/// A renderer that renders a "tile hunt".
///
/// A tile hunt marks visited tiles on a fixed zoom level.
///
/// The constructor takes the tile hunt level as parameter: TilehuntRenderer(zoom)
#[pyclass]
struct TilehuntRenderer {
    inner: renderer::tilehunt::Renderer,
}

#[pymethods]
impl TilehuntRenderer {
    #[new]
    fn new(zoom: u32) -> TilehuntRenderer {
        TilehuntRenderer {
            inner: renderer::tilehunt::Renderer::new(zoom),
        }
    }
}

/// 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>,
) -> PyResult<()> {
    let mut tracks = vec![];

    for item in items.try_iter()? {
        let item = item?;
        tracks.push(item.extract::<Track>()?.inner);
    }

    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 + Send),
) -> PyResult<()> {
    storage.prepare().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))?;

    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))?;

            storage.prepare_zoom(zoom).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))?;

        Ok(())
    })
}

/// Python bindings for the hittekaart heatmap tile generator.
///
/// hittekaart renders GPS tracks (usually from GPX files) to heatmap tiles for consumption with
/// OpenStreetMap data.
///
/// You can find more documentation about hittekaart in its README:
/// <https://gitlab.com/dunj3/hittekaart>
///
/// You can find rendered docs as part of fietsboek's documentation:
/// <https://docs.fietsboek.org/developer/module/hittekaart_py.html>
#[pymodule]
fn hittekaart_py(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<Track>()?;
    m.add_class::<HeatmapRenderer>()?;
    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("HitteError", py.get_type::<HitteError>())?;
    Ok(())
}