allow overriding jinja environment (#528)

nornir-automation · May 3, 2020 · 62fdd37 · 62fdd37
1 parent cbbcf83
commit 62fdd37
Showing 1 changed file with 16 additions and 5 deletions.
diff --git a/nornir/core/helpers/jinja_helper.py b/nornir/core/helpers/jinja_helper.py
@@ -7,22 +7,33 @@ def render_from_file(
     path: str,
     template: str,
     jinja_filters: Optional[Dict[str, Any]] = None,
+    jinja_env: Optional[Environment] = None,
     **kwargs: Any
 ) -> str:
     jinja_filters = jinja_filters or {}
-    env = Environment(
-        loader=FileSystemLoader(path), undefined=StrictUndefined, trim_blocks=True
-    )
+    if jinja_env:
+        env = jinja_env
+        env.loader = FileSystemLoader(path)
+    else:
+        env = Environment(
+            loader=FileSystemLoader(path), undefined=StrictUndefined, trim_blocks=True,
+        )
     env.filters.update(jinja_filters)
     t = env.get_template(template)
     return t.render(**kwargs)
 
 
 def render_from_string(
-    template: str, jinja_filters: Optional[Dict[str, Any]] = None, **kwargs: Any
+    template: str,
+    jinja_filters: Optional[Dict[str, Any]] = None,
+    jinja_env: Optional[Environment] = None,
+    **kwargs: Any
 ) -> str:
     jinja_filters = jinja_filters or {}
-    env = Environment(undefined=StrictUndefined, trim_blocks=True)
+    if jinja_env:
+        env = jinja_env
+    else:
+        env = Environment(undefined=StrictUndefined, trim_blocks=True,)
     env.filters.update(jinja_filters)
     t = env.from_string(template)
     return t.render(**kwargs)