feat: cache for jinja

Author: a3li

mapbuilder/builder.py

@@ -86,7 +86,7 @@ def __init__(self, source_dir, target_dir, cache_dir, config):
             else:
                 logging.error(f"Unknown data source type for data source {data_source}")
 
-        self.jinja_handler = JinjaHandler(self.data, self.config)
+        self.jinja_handler = JinjaHandler(self.data, self.config, cache_dir)
 
     def __load(self, name: str, src: str):
         if src.startswith("dfs:"):

mapbuilder/handlers/jinja.py

@@ -4,6 +4,7 @@
 import numpy as np
 import shapely
 import shapely.ops
+from cachelib import FileSystemCache
 from jinja2 import Environment, FileSystemLoader
 from more_itertools import unique_everseen
 from shapely import Geometry, Polygon
@@ -12,17 +13,20 @@
 from mapbuilder.utils.ad import render_cl, render_runways
 from mapbuilder.utils.ecl import draw_ecl_dashes, draw_loc_tick, draw_marker_ticks
 from mapbuilder.utils.geo import brg, fix
+from mapbuilder.utils.jinja_fragment_cache import FragmentCacheExtension
 from mapbuilder.utils.sidstar import render_sid
 
 
 class JinjaHandler:
-    def __init__(self, data, config):
+    def __init__(self, data, config, cache_dir):
         self.data = data
         self.config = config
+        self.cache_dir = cache_dir
 
     def handle(self, item: Path) -> str:
         file_loader = FileSystemLoader(item.parent)
-        jinja_env = Environment(loader=file_loader)
+        jinja_env = Environment(loader=file_loader, extensions=[FragmentCacheExtension])
+        jinja_env.fragment_cache = FileSystemCache(self.cache_dir)
         jinja_env.globals.update(
             data=self.data,
             runways=self.config["runways"],

mapbuilder/utils/jinja_fragment_cache.py

@@ -0,0 +1,54 @@
+from jinja2 import nodes
+from jinja2.ext import Extension
+
+
+class FragmentCacheExtension(Extension):
+    # a set of names that trigger the extension.
+    tags = {"cache"}
+
+    def __init__(self, environment):
+        super().__init__(environment)
+
+        # add the defaults to the environment
+        environment.extend(fragment_cache_prefix="", fragment_cache=None)
+
+    def parse(self, parser):
+        # the first token is the token that started the tag.  In our case
+        # we only listen to ``'cache'`` so this will be a name token with
+        # `cache` as value.  We get the line number so that we can give
+        # that line number to the nodes we create by hand.
+        lineno = next(parser.stream).lineno
+
+        # now we parse a single expression that is used as cache key.
+        args = [parser.parse_expression()]
+
+        # if there is a comma, the user provided a timeout.  If not use
+        # None as second parameter.
+        if parser.stream.skip_if("comma"):
+            args.append(parser.parse_expression())
+        else:
+            args.append(nodes.Const(None))
+
+        # now we parse the body of the cache block up to `endcache` and
+        # drop the needle (which would always be `endcache` in that case)
+        body = parser.parse_statements(["name:endcache"], drop_needle=True)
+
+        # now return a `CallBlock` node that calls our _cache_support
+        # helper method on this extension.
+        return nodes.CallBlock(
+            self.call_method("_cache_support", args), [], [], body
+        ).set_lineno(lineno)
+
+    def _cache_support(self, name, timeout, caller):
+        """Helper callback."""
+        key = self.environment.fragment_cache_prefix + name
+
+        # try to load the block from the cache
+        # if there is no fragment in the cache, render it and store
+        # it in the cache.
+        rv = self.environment.fragment_cache.get(key)
+        if rv is not None:
+            return rv
+        rv = caller()
+        self.environment.fragment_cache.add(key, rv, timeout)
+        return rv