From 22d02c007acfa1e4a2e011fe317d4696f3e1e6d6 Mon Sep 17 00:00:00 2001
From: Adrian Herrera <adrian.herrera02@gmail.com>
Date: Sun, 30 Dec 2018 14:02:26 +1100
Subject: [PATCH] templates: added ability to set custom template directory

Useful for custom projects outside of the s2e-env module.
---
 s2e_env/utils/templates.py | 41 +++++++++++++++++++++-----------------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/s2e_env/utils/templates.py b/s2e_env/utils/templates.py
index fbbc0e9c..11bd5779 100644
--- a/s2e_env/utils/templates.py
+++ b/s2e_env/utils/templates.py
@@ -29,7 +29,8 @@
 from .memoize import memoize
 
 
-TEMPLATES_DIR = os.path.join(os.path.dirname(__file__), '..', 'templates')
+DEFAULT_TEMPLATES_DIR = os.path.join(os.path.dirname(__file__), '..',
+                                     'templates')
 
 
 def _datetimefilter(value, format_='%H:%M %d-%m-%Y'):
@@ -40,11 +41,14 @@ def _datetimefilter(value, format_='%H:%M %d-%m-%Y'):
 
 
 @memoize
-def _init_template_env(templates_dir=TEMPLATES_DIR):
+def _init_template_env(templates_dir=None):
     """
     Initialize the jinja2 templating environment using the templates in the
     given directory.
     """
+    if not templates_dir:
+        templates_dir = DEFAULT_TEMPLATES_DIR
+
     env = Environment(loader=FileSystemLoader(templates_dir),
                       autoescape=False, undefined=StrictUndefined)
     env.filters['datetimefilter'] = _datetimefilter
@@ -52,29 +56,30 @@ def _init_template_env(templates_dir=TEMPLATES_DIR):
     return env
 
 
-def render_template(context, template, path=None, executable=False):
+def render_template(context, template, output_path=None, templates_dir=None,
+                    executable=False):
     """
     Renders the ``template`` template with the given ``context``. The result is
-    written to ``path``. If ``path`` is not specified, the result is
-    returned as a string
+    returned as a string and written to ``output_path`` (if specified).
+
+    A directory containing the Jinja templates can optionally be specified.
     """
-    env = _init_template_env()
-    data = env.get_template(template).render(context)
+    env = _init_template_env(templates_dir)
+
+    rendered_data = env.get_template(template).render(context)
 
     # Remove trailing spaces
     cleaned_lines = []
-    for line in data.splitlines():
+    for line in rendered_data.splitlines():
         cleaned_lines.append(line.rstrip())
-    data = '\n'.join(cleaned_lines)
-
-    if not path:
-        return data
+    rendered_data = '\n'.join(cleaned_lines)
 
-    with open(path, 'w') as f:
-        f.write(data)
+    if output_path:
+        with open(output_path, 'w') as f:
+            f.write(rendered_data)
 
-    if executable:
-        st = os.stat(path)
-        os.chmod(path, st.st_mode | stat.S_IEXEC)
+        if executable:
+            st = os.stat(output_path)
+            os.chmod(output_path, st.st_mode | stat.S_IEXEC)
 
-    return True
+    return rendered_data