diff --git a/aidge_core/export_utils/code_generation.py b/aidge_core/export_utils/code_generation.py index a02fc0966702cec7a2cbe33f8411bb71e3035e90..29b24bb82314cdd85db78c71b98bbead68a0b855 100644 --- a/aidge_core/export_utils/code_generation.py +++ b/aidge_core/export_utils/code_generation.py @@ -1,5 +1,5 @@ from pathlib import Path -from jinja2 import Environment, FileSystemLoader +from jinja2 import Environment, FileSystemLoader, StrictUndefined from typing import Union @@ -21,13 +21,9 @@ def generate_file(file_path: Union[Path, str], template_path: Union[Path, str], # Make dir file_path.parent.mkdir(parents=True, exist_ok=True) - # Select template - template = Environment(loader=FileSystemLoader( - template_path.parent)).get_template(template_path.name) - # Generate file with open(file_path, mode="w", encoding="utf-8") as file: - file.write(template.render(kwargs)) + file.write(generate_str(template_path, **kwargs)) def generate_str(template_path: Union[Path, str], **kwargs) -> str: @@ -43,4 +39,4 @@ def generate_str(template_path: Union[Path, str], **kwargs) -> str: if isinstance(template_path, str): template_path = Path(template_path) return Environment(loader=FileSystemLoader( - template_path.parent)).get_template(template_path.name).render(kwargs) + template_path.parent), undefined=StrictUndefined).get_template(template_path.name).render(kwargs)