Use a separate Jinja env per extension

Jinja Environments cache templates by name, therefore we cannot use
the same environment across extensions.

Signed-off-by: Andreas Gerstmayr <[email protected]>

src/fava/application.py

@@ -24,7 +24,6 @@
 from urllib.parse import urlparse
 from urllib.parse import urlunparse
 
-import jinja2
 import markdown2  # type: ignore[import]
 from beancount import __version__ as beancount_version
 from beancount.utils.text_utils import replace_numbers
@@ -322,16 +321,17 @@ def extension_js_module(extension_name: str) -> Response:
     @fava_app.route("/<bfile>/extension/<extension_name>/")
     def extension_report(extension_name: str) -> str:
         """Endpoint for extension reports."""
-        g.extension = g.ledger.extensions.get_extension(extension_name)
-        if not g.extension or g.extension.report_title is None:
+        ext = g.ledger.extensions.get_extension(extension_name)
+        if ext is None or ext.report_title is None:
             abort(404)
-        content = Markup(
-            render_template(f"{g.extension.name}.html", extension=g.extension),
-        )
+
+        g.extension = ext
+        template = ext.jinja_env.get_template(f"{ext.name}.html")
+        content = Markup(template.render(ledger=g.ledger, extension=ext))
         return render_template(
             "_layout.html",
             content=content,
-            page_title=g.extension.report_title,
+            page_title=ext.report_title,
         )
 
     @fava_app.route("/<bfile>/download-query/query_result.<result_format>")
@@ -421,18 +421,6 @@ def _get_locale() -> str | None:
         babel = Babel(fava_app, locale_selector=_get_locale)
 
 
-def _ext_template_loader_func(name: str) -> str | None:
-    """Load a template from the extension folder if an extension is active."""
-    if g.extension:
-        template_path = g.extension.extension_dir / "templates" / name
-        if template_path.exists():
-            return template_path.read_text(encoding="utf-8")
-    return None
-
-
-_ext_template_loader = jinja2.FunctionLoader(_ext_template_loader_func)
-
-
 def create_app(
     files: Iterable[Path | str],
     load: bool = False,
@@ -454,12 +442,6 @@ def create_app(
     _setup_babel(fava_app)
     _setup_filters(fava_app, read_only=read_only, incognito=incognito)
     _setup_routes(fava_app)
-    loader = fava_app.jinja_loader
-    if loader is None:
-        raise ValueError("Expected Flask app to have jinja_loader.")
-    fava_app.jinja_loader = jinja2.ChoiceLoader(  # type: ignore[assignment]
-        [loader, _ext_template_loader],
-    )
 
     fava_app.config["HAVE_EXCEL"] = HAVE_EXCEL
     fava_app.config["BEANCOUNT_FILES"] = [str(f) for f in files]

src/fava/ext/__init__.py

@@ -5,10 +5,14 @@
 import importlib
 import inspect
 import sys
+from functools import cached_property
 from pathlib import Path
 from typing import Any
 from typing import TYPE_CHECKING
 
+import jinja2
+from flask import current_app
+
 from fava.helpers import BeancountError
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -59,6 +63,15 @@ def extension_dir(self) -> Path:
         """Directory to look for templates directory and Javascript code."""
         return Path(inspect.getfile(self.__class__)).parent
 
+    @cached_property
+    def jinja_env(self) -> jinja2.Environment:
+        """Jinja env for this extension."""
+        if not current_app.jinja_loader:
+            raise ValueError("Expected Flask app to have jinja_loader.")
+        ext_loader = jinja2.FileSystemLoader(self.extension_dir / "templates")
+        loader = jinja2.ChoiceLoader([ext_loader, current_app.jinja_loader])
+        return current_app.jinja_env.overlay(loader=loader)
+
     def after_entry_modified(self, entry: Directive, new_lines: str) -> None:
         """Run after an `entry` has been modified."""