From e3338a1f184533f0363a1f8ded15948ee15827cb Mon Sep 17 00:00:00 2001
From: Konstantin Demin <rockdrilla@gmail.com>
Date: Mon, 15 Jul 2024 21:51:40 +0300
Subject: [PATCH] rewrite jinja.py

---
 scripts/jinja.py | 77 ++++++++++++++++++++++++++++++------------------
 1 file changed, 49 insertions(+), 28 deletions(-)

diff --git a/scripts/jinja.py b/scripts/jinja.py
index 26ea955..a592136 100755
--- a/scripts/jinja.py
+++ b/scripts/jinja.py
@@ -8,24 +8,42 @@ import sys
 import jinja2
 import yaml
 
-NGX_JINJA_MODULES = sorted(set(os.getenv('NGX_JINJA_MODULES', 'os os.path sys netaddr psutil re wcmatch').split(sep=' ')))
 
-ME = sys.argv[0]
-
-NGX_JINJA_SUFFIX = os.getenv('NGX_JINJA_SUFFIX', '.j2')
-if NGX_JINJA_SUFFIX == '':
-    raise ValueError('NGX_JINJA_SUFFIX is empty')
-if not NGX_JINJA_SUFFIX.startswith('.'):
-    raise ValueError('NGX_JINJA_SUFFIX does not start with dot (".")')
-
-NGX_JINJA_CONFIG = os.getenv('NGX_JINJA_CONFIG', '')
-if (NGX_JINJA_CONFIG != '') and (not os.path.exists(NGX_JINJA_CONFIG)):
-    print(f'{ME}: config does not exist, skipping: {NGX_JINJA_CONFIG}', file=sys.stderr)
-
-NGX_JINJA_CONFIG_NAMES = [
+J2_MODULES_DEFAULT = 'os os.path sys netaddr psutil re wcmatch'
+J2_SUFFIX_DEFAULT = '.j2'
+J2_CFG_PATHS = [
     '/angie/jinja',
     '/etc/angie/jinja',
 ]
+J2_CFG_EXTS = [
+    'yml',
+    'yaml',
+    'json',
+]
+J2_SEARCH_PATH = [
+    '.',
+    '/etc/angie',
+]
+J2_EXT_LIST = [
+    'jinja2.ext.do',
+    'jinja2.ext.loopcontrols',
+]
+
+
+J2_MODULES = sorted(set(os.getenv('NGX_JINJA_MODULES', J2_MODULES_DEFAULT).split(sep=' ')))
+
+ME = sys.argv[0]
+
+J2_SUFFIX = os.getenv('NGX_JINJA_SUFFIX', J2_SUFFIX_DEFAULT)
+if J2_SUFFIX == '':
+    raise ValueError('NGX_JINJA_SUFFIX is empty')
+if not J2_SUFFIX.startswith('.'):
+    raise ValueError('NGX_JINJA_SUFFIX does not start with dot (".")')
+
+J2_CONFIG = os.getenv('NGX_JINJA_CONFIG', '')
+if (J2_CONFIG != '') and (not os.path.exists(J2_CONFIG)):
+    print(f'{ME}: config does not exist, skipping: {J2_CONFIG}', file=sys.stderr)
+
 
 if len(sys.argv) < 2:
     raise ValueError('not enough arguments (needed: 2)')
@@ -44,14 +62,14 @@ if len(sys.argv) == 3:
     output_file = sys.argv[2]
 else:
     output_file, ext = os.path.splitext(input_file)
-    if ext != NGX_JINJA_SUFFIX:
-        raise ValueError(f'input file name extension mismatch (not a "{NGX_JINJA_SUFFIX}")')
+    if ext != J2_SUFFIX:
+        raise ValueError(f'input file name extension mismatch (not a "{J2_SUFFIX}")')
 
 if input_file == output_file:
     raise ValueError('unable to process template inplace')
 
 kwargs = {}
-for m in NGX_JINJA_MODULES:
+for m in J2_MODULES:
     kwargs[m] = importlib.import_module(m)
 kwargs['env'] = os.environ
 kwargs['cfg'] = {}
@@ -83,16 +101,19 @@ def merge_dict_from_file(filename):
     return False
 
 
-for base_name in NGX_JINJA_CONFIG_NAMES:
-    for file_name in [ base_name + '.' + ext for ext in [ 'yml', 'yaml', 'json' ] ]:
-        if merge_dict_from_file(file_name):
-            break
-        continue
-merge_dict_from_file(NGX_JINJA_CONFIG)
+if (J2_CONFIG != '') and (os.path.isfile(J2_CONFIG)):
+    merge_dict_from_file(J2_CONFIG)
+else:
+    for base in J2_CFG_PATHS:
+        for full in [ base + '.' + ext for ext in J2_CFG_EXTS ]:
+            if merge_dict_from_file(full):
+                break
+            continue
 
-with open(input_file, mode='r', encoding='utf-8') as fin:
-    template = jinja2.Template(fin.read())
-rendered = template.render(**kwargs)
 
-with open(output_file, mode='w', encoding='utf-8') as fout:
-    fout.write(rendered)
+j2_loader = jinja2.FileSystemLoader(J2_SEARCH_PATH, followlinks=True)
+j2_environ = jinja2.Environment(loader=j2_loader, extensions=J2_EXT_LIST)
+j2_template = j2_environ.get_template(input_file)
+j2_stream = j2_template.stream(**kwargs)
+j2_stream.disable_buffering()
+j2_stream.dump(output_file, encoding='utf-8')