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(),
        }
    }
}

#[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()))
    }
}

impl Storage {
    fn open(&self) -> PyResult<Box<dyn hittekaart::storage::Storage>> {
        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,
        }
    }
}

/// Generate the heatmap.
///
/// * items is an iterable of Track
/// * renderer should be one of the renderers (such as HeatmapRenderer)
/// * storage is the Storage output
#[pyfunction]
fn generate(
    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);
    }

    if let Ok(r) = renderer.downcast::<HeatmapRenderer>() {
        do_generate(tracks, &r.borrow().inner, &mut *storage.borrow().open()?)
    } else {
        Err(PyTypeError::new_err("Expected a HeatmapRenderer"))
    }
}

fn do_generate<R: Renderer>(
    tracks: Vec<Vec<Coordinates>>,
    renderer: &R,
    storage: &mut dyn hittekaart::storage::Storage,
) -> 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))?;

        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(())
}

/// 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))
}

/// 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::<Storage>()?;
    m.add_function(wrap_pyfunction!(generate, m)?)?;
    m.add_function(wrap_pyfunction!(set_threads, m)?)?;
    m.add("HitteError", py.get_type::<HitteError>())?;
    Ok(())
}