aos: Clean up the jinja2_template rule a little bit

I made some enhancements for custom filters and generated
replacements. It was easier to implement those after cleaning up the
code a bit first. This patch does that.

Change-Id: Idbdfae27cf1a3bec598eeaeabf7e7b9e36341242
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 82222c3..047a805 100644
--- a/tools/build_rules/jinja2_generator.py
+++ b/tools/build_rules/jinja2_generator.py
@@ -3,6 +3,7 @@
 import argparse
 import json
 import sys
+from pathlib import Path
 
 import jinja2
 
@@ -13,25 +14,27 @@
     # a template.
     parser = argparse.ArgumentParser(
         description="Generates the raspberry pi configs from a template.")
-    parser.add_argument("template", type=str, help="File to use for template.")
+    parser.add_argument("template",
+                        type=Path,
+                        help="File to use for template.")
     parser.add_argument(
         "replacements",
         type=json.loads,
         help="Dictionary of parameters to replace in the template.")
-    parser.add_argument("output", type=str, help="Output file to create.")
     parser.add_argument(
-        "genfiles_dir",
-        type=str,
-        help="Directory where generated JSON files will be available.")
+        "--include_dir",
+        action="append",
+        type=Path,
+        default=[],
+        help="One or more search directories for {% include %} blocks.",
+    )
+    parser.add_argument("output", type=Path, help="Output file to create.")
     args = parser.parse_args(sys.argv[1:])
 
-    with open(args.template, 'r') as input_file:
-        template = jinja2.Environment(loader=jinja2.FileSystemLoader(
-            [".", args.genfiles_dir])).from_string(input_file.read())
+    env = jinja2.Environment(loader=jinja2.FileSystemLoader(args.include_dir))
+    template = env.from_string(args.template.read_text())
 
-    output = template.render(args.replacements)
-    with open(args.output, 'w') as config_file:
-        config_file.write(output)
+    args.output.write_text(template.render(args.replacements))
 
 
 if __name__ == '__main__':
diff --git a/tools/build_rules/template.bzl b/tools/build_rules/template.bzl
index d3547d8..28d3c82 100644
--- a/tools/build_rules/template.bzl
+++ b/tools/build_rules/template.bzl
@@ -3,13 +3,28 @@
     parameters = dict(ctx.attr.parameters)
     parameters.update(ctx.attr.list_parameters)
 
-    ctx.actions.run_shell(
+    # For now we don't really want the user to worry about which configuration
+    # to pull the file from. We don't yet have a use case for pulling the same
+    # file from multiple configurations. We point Jinja at all the configuration
+    # roots.
+    include_dirs = depset([
+        file.root.path or "."
+        for file in ctx.files.includes
+    ]).to_list()
+
+    args = ctx.actions.args()
+    args.add(ctx.file.src)
+    args.add(json.encode(parameters))
+    args.add(out)
+    args.add_all(include_dirs, before_each = "--include_dir")
+
+    ctx.actions.run(
         inputs = ctx.files.src + ctx.files.includes,
         tools = [ctx.executable._jinja2],
         progress_message = "Generating " + out.short_path,
         outputs = [out],
-        # TODO(james): Is the genfiles_dir the correct thing?
-        command = ctx.executable._jinja2.path + " " + ctx.files.src[0].path + " '" + str(parameters) + "' " + out.path + " " + ctx.genfiles_dir.path,
+        executable = ctx.executable._jinja2,
+        arguments = [args],
     )
 
     return [DefaultInfo(files = depset([out])), OutputGroupInfo(out = depset([out]))]
@@ -54,4 +69,11 @@
     # differ from `out`, name the rule as the `name` plus a suffix
     rule_name = name + "_rule"
 
-    jinja2_template_rule(name = rule_name, out = name, src = src, parameters = parameters, list_parameters = list_parameters, **kwargs)
+    jinja2_template_rule(
+        name = rule_name,
+        out = name,
+        src = src,
+        parameters = parameters,
+        list_parameters = list_parameters,
+        **kwargs
+    )