diff options
| -rw-r--r-- | fietsboek/config.py | 3 | ||||
| -rw-r--r-- | fietsboek/data.py | 11 | ||||
| -rw-r--r-- | fietsboek/hittekaart.py | 125 | ||||
| -rw-r--r-- | fietsboek/scripts/fietsctl.py | 48 |
4 files changed, 186 insertions, 1 deletions
diff --git a/fietsboek/config.py b/fietsboek/config.py index a5a87a8..3fc191f 100644 --- a/fietsboek/config.py +++ b/fietsboek/config.py @@ -195,6 +195,9 @@ class Config(BaseModel): tile_layers: typing.List[TileLayerConfig] = [] """Tile layers.""" + hittekaart_bin: str = Field("", alias="hittekaart.bin") + """Path to the hittekaart binary.""" + @validator("session_key") def _good_session_key(cls, value): """Ensures that the session key has been changed from its default diff --git a/fietsboek/data.py b/fietsboek/data.py index 9be47da..7457986 100644 --- a/fietsboek/data.py +++ b/fietsboek/data.py @@ -82,6 +82,17 @@ class DataManager: path.mkdir(parents=True) return TrackDataDir(track_id, path, journal=True, is_fresh=True) + def initialize_user(self, user_id: int) -> "UserDataDir": + """Creates the data directory for a user. + + :raises FileExistsError: If the directory already exists. + :param user_id: ID of the user. + :return: The manager that can be used to manage this user's data. + """ + path = self._user_data_dir(user_id) + path.mkdir(parents=True) + return UserDataDir(user_id, path) + def purge(self, track_id: int): """Forcefully purges all data from the given track. diff --git a/fietsboek/hittekaart.py b/fietsboek/hittekaart.py new file mode 100644 index 0000000..f580c51 --- /dev/null +++ b/fietsboek/hittekaart.py @@ -0,0 +1,125 @@ +"""Interface to the hittekaart_ application to generate heatmaps. + +.. _hittekaart: https://gitlab.com/dunj3/hittekaart +""" +import enum +import logging +import shutil +import subprocess +import tempfile +from pathlib import Path +from typing import Optional + +from sqlalchemy import select +from sqlalchemy.orm import aliased +from sqlalchemy.orm.session import Session + +from . import models +from .data import DataManager +from .models.track import TrackType + +LOGGER = logging.getLogger(__name__) + + +class Mode(enum.Enum): + """Heatmap generation mode. + + This enum represents the different types of overlay maps that + ``hittekaart`` can generate. + """ + + HEATMAP = "heatmap" + TILEHUNTER = "tilehunter" + + +def generate( + output: Path, + mode: Mode, + input_files: list[Path], + *, + exe_path: Optional[Path] = None, + threads: int = 0, +): + """Calls hittekaart with the given arguments. + + :param output: Output filename. Note that this function always uses the + sqlite output mode. + :param mode: What to generate. + :param input_files: List of paths to the input files. + :param exe_path: Path to the hittekaart binary. If not given, + ``hittekaart`` is searched in the path. + :param threads: Number of threads that ``hittekaart`` should use. Defaults + to 0, which uses all available cores. + """ + # There are two reasons why we do the tempfile dance: + # 1. hittekaart refuses to overwrite existing files + # 2. This way we can (hope for?) an atomic move (at least if temporary file + # is on the same filesystem). In the future, we might want to enforce + # this, but for now, it's alright. + with tempfile.TemporaryDirectory() as tempdir: + tmpfile = Path(tempdir) / "hittekaart.sqlite" + binary = str(exe_path) if exe_path else "hittekaart" + cmdline = [ + binary, + "--sqlite", + "-o", + str(tmpfile), + "-m", + mode.value, + "-t", + str(threads), + "--", + ] + cmdline.extend(map(str, input_files)) + LOGGER.debug("Running %r", cmdline) + subprocess.run(cmdline, check=True) + + LOGGER.debug("Moving temporary file") + shutil.move(tmpfile, output) + + +def generate_for( + user: models.User, + dbsession: Session, + data_manager: DataManager, + mode: Mode, + *, + exe_path: Optional[Path] = None, + threads: int = 0, +): + """Uses :meth:`generate` to generate a heatmap for the given user. + + This function automatically retrieves the user's tracks from the database + and passes them to ``hittekaart``. + + The output is saved in the user's data directory using the + ``data_manager``. + + :param user: The user for which to generate the map. + :param dbsession: The database session. + :param data_manager: The data manager. + :param mode: The mode of the heatmap. + :param exe_path: See :meth:`generate`. + :param threads: See :meth:`generate`. + """ + query = user.all_tracks_query() + query = select(aliased(models.Track, query)).where(query.c.type == TrackType.ORGANIC) + input_paths = [] + for track in dbsession.execute(query).scalars(): + path = data_manager.open(track.id).gpx_path() + input_paths.append(path) + + try: + user_dir = data_manager.initialize_user(user.id) + except FileExistsError: + user_dir = data_manager.open_user(user.id) + + output_paths = { + Mode.HEATMAP: user_dir.heatmap_path(), + Mode.TILEHUNTER: user_dir.tilehunt_path(), + } + + generate(output_paths[mode], mode, input_paths, exe_path=exe_path, threads=threads) + + +__all__ = ["Mode", "generate", "generate_for"] diff --git a/fietsboek/scripts/fietsctl.py b/fietsboek/scripts/fietsctl.py index 0043862..8be226a 100644 --- a/fietsboek/scripts/fietsctl.py +++ b/fietsboek/scripts/fietsctl.py @@ -8,7 +8,7 @@ from pyramid.paster import bootstrap, setup_logging from pyramid.scripting import AppEnvironment from sqlalchemy import select -from .. import __VERSION__, models +from .. import __VERSION__, hittekaart, models from . import config_option EXIT_OKAY = 0 @@ -219,6 +219,52 @@ def cmd_maintenance_mode(ctx: click.Context, config: str, disable: bool, reason: (data_manager.data_dir / "MAINTENANCE").write_text(reason, encoding="utf-8") +@cli.command("hittekaart") +@config_option +@click.option( + "--mode", + "modes", + help="Heatmap type to generate", + type=click.Choice([mode.value for mode in hittekaart.Mode]), + multiple=True, + default=["heatmap"], +) +@optgroup.group("User selection", cls=RequiredMutuallyExclusiveOptionGroup) +@optgroup.option("--id", "-i", "id_", help="database ID of the user", type=int) +@optgroup.option("--email", "-e", help="email of the user") +@click.pass_context +def cmd_hittekaart( + ctx: click.Context, + config: str, + modes: list[str], + id_: Optional[int], + email: Optional[str], +): + """Generate heatmap for a user.""" + env = setup(config) + modes = [hittekaart.Mode(mode) for mode in modes] + + if id_ is not None: + query = select(models.User).filter_by(id=id_) + else: + query = models.User.query_by_email(email) + + exe_path = env["request"].config.hittekaart_bin + with env["request"].tm: + dbsession = env["request"].dbsession + data_manager = env["request"].data_manager + user = dbsession.execute(query).scalar_one_or_none() + if user is None: + click.echo("Error: No such user found.", err=True) + ctx.exit(EXIT_FAILURE) + + click.echo(f"Generating overlay maps for {user.name}...") + + for mode in modes: + hittekaart.generate_for(user, dbsession, data_manager, mode, exe_path=exe_path) + click.echo(f"Generated {mode.value}") + + @cli.command("version") def cmd_version(): """Show the installed fietsboek version.""" |