move to hittekaart-py instead of subprocess'ing

5 files changed, 56 insertions, 30 deletions

diff --git a/fietsboek/hittekaart.py b/fietsboek/hittekaart.py
index 15f2855..06319a4 100644
--- a/fietsboek/hittekaart.py
+++ b/fietsboek/hittekaart.py
@@ -6,11 +6,11 @@
 import enum
 import logging
 import shutil
-import subprocess
 import tempfile
 from pathlib import Path
 from typing import Optional
 
+import hittekaart_py
 from sqlalchemy import select
 from sqlalchemy.orm import aliased
 from sqlalchemy.orm.session import Session
@@ -21,6 +21,13 @@ from .models.track import TrackType
 
 LOGGER = logging.getLogger(__name__)
 
+COMPRESSION_MAP = {
+    ".br": "brotli",
+    ".gz": "gzip",
+}
+
+TILEHUNTER_ZOOM = 14
+
 
 class Mode(enum.Enum):
     """Heatmap generation mode.
@@ -38,7 +45,6 @@ def generate(
     mode: Mode,
     input_files: list[Path],
     *,
-    exe_path: Optional[Path] = None,
     threads: int = 0,
 ):
     """Calls hittekaart with the given arguments.
@@ -47,13 +53,23 @@ def generate(
         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.
     """
     if not input_files:
         return
+
+    if mode == Mode.HEATMAP:
+        renderer = hittekaart_py.HeatmapRenderer()
+    elif mode == Mode.TILEHUNTER:
+        renderer = hittekaart_py.TilehuntRenderer(TILEHUNTER_ZOOM)
+
+    LOGGER.debug("Loading tracks ...")
+    tracks = [
+        hittekaart_py.Track.from_file(bytes(input_file), COMPRESSION_MAP.get(input_file.suffix))
+        for input_file in input_files
+    ]
+    LOGGER.debug("Tracks loaded!")
     # 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
@@ -61,23 +77,12 @@ def generate(
     #    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, stdout=subprocess.DEVNULL)
-
-        LOGGER.debug("Moving temporary file")
+        sink = hittekaart_py.Storage.Sqlite(bytes(tmpfile))
+        settings = hittekaart_py.Settings(threads=threads)
+        LOGGER.debug("Running hittekaart (renderer %r) to %r", renderer, tmpfile)
+        hittekaart_py.generate(settings, tracks, renderer, sink)
+
+        LOGGER.debug("Moving temporary file %r to %r", tmpfile, output)
         shutil.move(tmpfile, output)
 
 
@@ -87,7 +92,6 @@ def generate_for(
     data_manager: DataManager,
     mode: Mode,
     *,
-    exe_path: Optional[Path] = None,
     threads: int = 0,
 ):
     """Uses :meth:`generate` to generate a heatmap for the given user.
@@ -104,7 +108,6 @@ def generate_for(
     :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`.
     """
     # pylint: disable=too-many-arguments
@@ -133,7 +136,7 @@ def generate_for(
         Mode.TILEHUNTER: user_dir.tilehunt_path(),
     }
 
-    generate(output_paths[mode], mode, input_paths, exe_path=exe_path, threads=threads)
+    generate(output_paths[mode], mode, input_paths, threads=threads)
 
 
 __all__ = ["Mode", "generate", "generate_for"]
diff --git a/fietsboek/scripts/fietscron.py b/fietsboek/scripts/fietscron.py
index a3f3f54..75575d8 100644
--- a/fietsboek/scripts/fietscron.py
+++ b/fietsboek/scripts/fietscron.py
@@ -142,7 +142,6 @@ def run_hittekaart(engine: Engine, data_manager: DataManager, redis: Redis, conf
     # re-generate all maps over time (e.g. if the hittekaart version changes or
     # we miss an update).
     modes = [hittekaart.Mode(mode) for mode in config.hittekaart_autogenerate]
-    exe_path = Path(config.hittekaart_bin) if config.hittekaart_bin else None
     session = Session(engine)
     had_hq_item = False
 
@@ -164,7 +163,6 @@ def run_hittekaart(engine: Engine, data_manager: DataManager, redis: Redis, conf
                 session,
                 data_manager,
                 mode,
-                exe_path=exe_path,
                 threads=config.hittekaart_threads,
             )
 
@@ -188,7 +186,7 @@ def run_hittekaart(engine: Engine, data_manager: DataManager, redis: Redis, conf
     for mode in modes:
         LOGGER.info("Generating %s for user %d (low-priority)", mode.value, user.id)
         hittekaart.generate_for(
-            user, session, data_manager, mode, exe_path=exe_path, threads=config.hittekaart_threads
+            user, session, data_manager, mode, threads=config.hittekaart_threads
         )
 
 
diff --git a/fietsboek/scripts/fietsctl.py b/fietsboek/scripts/fietsctl.py
index ad06e91..87d49e8 100644
--- a/fietsboek/scripts/fietsctl.py
+++ b/fietsboek/scripts/fietsctl.py
@@ -313,7 +313,6 @@ def cmd_user_hittekaart(
     else:
         query = models.User.query_by_email(email)
 
-    exe_path = env["request"].config.hittekaart_bin
     threads = env["request"].config.hittekaart_threads
     with env["request"].tm:
         dbsession = env["request"].dbsession
@@ -338,7 +337,7 @@ def cmd_user_hittekaart(
 
         for mode in modes:
             hittekaart.generate_for(
-                user, dbsession, data_manager, mode, exe_path=exe_path, threads=threads
+                user, dbsession, data_manager, mode, threads=threads
             )
             click.echo(f"Generated {mode.value}")
 
diff --git a/poetry.lock b/poetry.lock
index 440a4b2..c167492 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -788,6 +788,24 @@ docs = ["Sphinx", "furo"]
 test = ["objgraph", "psutil"]
 
 [[package]]
+name = "hittekaart-py"
+version = "0.1.0"
+description = ""
+optional = true
+python-versions = ">=3.8"
+groups = ["main"]
+markers = "extra == \"hittekaart\""
+files = []
+develop = false
+
+[package.source]
+type = "git"
+url = "https://gitlab.com/dunj3/hittekaart"
+reference = "6da6f0b5ce4cfb41aaf79699e289be678007e2ad"
+resolved_reference = "6da6f0b5ce4cfb41aaf79699e289be678007e2ad"
+subdirectory = "hittekaart-py"
+
+[[package]]
 name = "hupper"
 version = "1.12.1"
 description = "Integrated process monitor for developing and reloading daemons."
@@ -2562,7 +2580,10 @@ transaction = ">=1.6.0"
 [package.extras]
 test = ["zope.testing"]
 
+[extras]
+hittekaart = ["hittekaart-py"]
+
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.11"
-content-hash = "e35a5cd2b50b9643fbca57a4c55384041263bd8fffb5b53f118b5040c3fa2b83"
+content-hash = "d455ec6a90988d2e6567d911e3b332162671e25b1335930413f926e6aa60d26a"
diff --git a/pyproject.toml b/pyproject.toml
index be56bfa..c47d037 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,6 +49,11 @@ documentation = "https://docs.fietsboek.org/"
 homepage = "https://fietsboek.org/"
 repository = "https://gitlab.com/dunj3/fietsboek"
 
+[project.optional-dependencies]
+hittekaart = [
+    "hittekaart-py @ git+https://gitlab.com/dunj3/hittekaart@6da6f0b5ce4cfb41aaf79699e289be678007e2ad#subdirectory=hittekaart-py",
+]
+
 [tool.poetry]
 classifiers = [
     'Development Status :: 3 - Alpha',