From 4d9346a3657ec7b48fd59b6a16c1d7913add0664 Mon Sep 17 00:00:00 2001
From: Liam Beguin <liambeguin@gmail.com>
Date: Sun, 9 Jun 2024 17:51:28 -0400
Subject: [PATCH] Add jinja2 preprocessor stage

This allows users to leverage jinja to include other files but also
loops, conditions, and other templating features.

Signed-off-by: Liam Beguin <liambeguin@gmail.com>
---
 examples/jinja/connectors.yml |  7 +++++++
 examples/jinja/top.yml        | 21 +++++++++++++++++++++
 setup.py                      |  1 +
 src/wireviz/wv_cli.py         | 25 +++++++++++++++++++++++--
 4 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 examples/jinja/connectors.yml
 create mode 100644 examples/jinja/top.yml

diff --git a/examples/jinja/connectors.yml b/examples/jinja/connectors.yml
new file mode 100644
index 00000000..9f5a6ddc
--- /dev/null
+++ b/examples/jinja/connectors.yml
@@ -0,0 +1,7 @@
+connectors:
+{%- for i in range(1, 5) %}
+  X{{i}}:
+    type: Molex KK 254
+    subtype: female
+    pinlabels: [GND, RX, TX]
+{%- endfor %}
diff --git a/examples/jinja/top.yml b/examples/jinja/top.yml
new file mode 100644
index 00000000..b5602333
--- /dev/null
+++ b/examples/jinja/top.yml
@@ -0,0 +1,21 @@
+{% include 'connectors.yml' %}
+
+cables:
+{%- for i in range(1, 3) %}
+  W{{i}}:
+    gauge: 0.25 mm2
+    length: 0.2
+    color_code: DIN
+    wirecount: 3
+    shield: true
+{%- endfor %}
+
+connections:
+  -
+    - X1: [1,2,3]
+    - W1: [1,2,3]
+    - X2: [1,2,3]
+  -
+    - X3: [1,2,3]
+    - W2: [1,2,3]
+    - X4: [1,2,3]
diff --git a/setup.py b/setup.py
index 8cdd4864..2a036a83 100644
--- a/setup.py
+++ b/setup.py
@@ -22,6 +22,7 @@
         "pyyaml",
         "pillow",
         "graphviz",
+        "jinja2",
     ],
     license="GPLv3",
     keywords="cable connector hardware harness wiring wiring-diagram wiring-harness",
diff --git a/src/wireviz/wv_cli.py b/src/wireviz/wv_cli.py
index afb02494..7967eb78 100644
--- a/src/wireviz/wv_cli.py
+++ b/src/wireviz/wv_cli.py
@@ -3,6 +3,7 @@
 import os
 import sys
 from pathlib import Path
+import jinja2
 
 import click
 
@@ -50,6 +51,13 @@
     type=Path,
     help="YAML file to prepend to the input file (optional).",
 )
+@click.option(
+    "-I",
+    "--include-path",
+    default=None,
+    type=Path,
+    help="Include path used for Jinja2 templates",
+)
 @click.option(
     "-o",
     "--output-dir",
@@ -71,7 +79,7 @@
     default=False,
     help=f"Output {APP_NAME} version and exit.",
 )
-def wireviz(file, format, prepend, output_dir, output_name, version):
+def wireviz(file, format, prepend, include_path, output_dir, output_name, version):
     """
     Parses the provided FILE and generates the specified outputs.
     """
@@ -115,6 +123,15 @@ def wireviz(file, format, prepend, output_dir, output_name, version):
     else:
         prepend_input = ""
 
+
+    searchpath = [Path(f).parent for f in filepaths]
+    if include_path is not None:
+        searchpath.append(include_path)
+
+    env = jinja2.Environment(
+        loader=jinja2.FileSystemLoader(searchpath=searchpath),
+    )
+
     # run WireVIz on each input file
     for file in filepaths:
         file = Path(file)
@@ -130,7 +147,11 @@ def wireviz(file, format, prepend, output_dir, output_name, version):
             "Output file: ", f"{Path(_output_dir / _output_name)}.{output_formats_str}"
         )
 
-        yaml_input = open_file_read(file).read()
+        template = env.get_template(file.name)
+        yaml_input = template.render()
+        with open(Path(_output_dir / (_output_name + '.rendered.yml')), 'w') as f:
+            f.write(yaml_input)
+
         file_dir = file.parent
 
         yaml_input = prepend_input + yaml_input