aos: Support custom filters in jinja2_template

I've been wanting to do this a few times now and I decided it was
finally time to do it.

I'm looking to render some jinja2 templates where I want to turn an
array of numbers (RGB colors) into an object initialization. That's a
little cumbersome for the existing filters. It's much easier to do in
a custom filter. This patch enables custom filters.

Change-Id: I09317cba7ec916cbbf076924c2d33509c20a5889
Signed-off-by: James Kuszmaul <james.kuszmaul@bluerivertech.com>
diff --git a/tools/build_rules/jinja2_generator.py b/tools/build_rules/jinja2_generator.py
index 7bb9c75..30432ea 100644
--- a/tools/build_rules/jinja2_generator.py
+++ b/tools/build_rules/jinja2_generator.py
@@ -1,6 +1,7 @@
 #!/usr/bin/python3
 
 import argparse
+import importlib.util
 import json
 import sys
 from pathlib import Path
@@ -8,6 +9,31 @@
 import jinja2
 
 
+def load_filter_file(filename: Path, env: jinja2.Environment):
+    """Adds filters specified in the .py file.
+
+    The .py file has to define a `register_filters` function that will be
+    invoked. The function will be passed the jinja2 environment.
+
+        def register_filters(env: jinja2.Environment):
+            env.filters["custom_filter"] = ...
+
+    Then you can use it in the template.
+
+        Hello {{ "world" | custom_filter }}!
+
+    Based on https://stackoverflow.com/a/51575312.
+
+    Args:
+        filename: The .py file to load.
+        env: The environment to pass to the `register_filters` function.
+    """
+    spec = importlib.util.spec_from_file_location("filter_module", filename)
+    filter_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(filter_module)
+    filter_module.register_filters(env)
+
+
 def main():
     # Note: this is a pretty transparent interface to jinja2--there's no reason
     # this script couldn't be renamed and then used to generate any config from
@@ -36,9 +62,20 @@
         help="One or more search directories for {% include %} blocks.",
     )
     parser.add_argument("output", type=Path, help="Output file to create.")
+    parser.add_argument(
+        "--filter_file",
+        action="append",
+        type=str,
+        default=[],
+        help=("A .py file with a register_filters() function for custom "
+              "jinja2 filters."),
+    )
     args = parser.parse_args(sys.argv[1:])
 
     env = jinja2.Environment(loader=jinja2.FileSystemLoader(args.include_dir))
+    for filename in args.filter_file:
+        load_filter_file(filename, env)
+
     template = env.from_string(args.template.read_text())
 
     replacements = args.replacements.copy()
diff --git a/tools/build_rules/template.bzl b/tools/build_rules/template.bzl
index 7174969..b8b513b 100644
--- a/tools/build_rules/template.bzl
+++ b/tools/build_rules/template.bzl
@@ -19,9 +19,10 @@
     args.add_all(include_dirs, before_each = "--include_dir")
     if ctx.file.parameters_file:
         args.add("--replacements_file", ctx.file.parameters_file)
+    args.add_all(ctx.files.filter_srcs, before_each = "--filter_file")
 
     ctx.actions.run(
-        inputs = ctx.files.src + ctx.files.includes + ctx.files.parameters_file,
+        inputs = ctx.files.src + ctx.files.includes + ctx.files.parameters_file + ctx.files.filter_srcs,
         tools = [ctx.executable._jinja2],
         progress_message = "Generating " + out.short_path,
         outputs = [out],
@@ -60,6 +61,11 @@
             allow_files = True,
             doc = """Files which are included by the template.""",
         ),
+        "filter_srcs": attr.label_list(
+            allow_files = [".py"],
+            doc = """Files that are sourced for filters.
+Needs to have a register_filters function defined.""",
+        ),
         "_jinja2": attr.label(
             default = "//tools/build_rules:jinja2_generator",
             cfg = "exec",