Make it flexible to load macros in SSG

ssg/jinja.py

@@ -6,6 +6,7 @@
 from __future__ import print_function
 
 import os.path
+import sys
 import jinja2
 
 try:
@@ -208,41 +209,92 @@ def add_python_functions(substitutions_dict):
     substitutions_dict['expand_yaml_path'] = expand_yaml_path
 
 
-def load_macros(substitutions_dict=None):
+def _load_macros_from_directory(macros_directory, substitutions_dict):
     """
-    Augments the provided substitutions_dict with project Jinja macros found in the /shared/ directory.
-
-    This function loads Jinja macro files from a predefined directory, processes them, and updates
-    the substitutions_dict with the macro definitions. If no substitutions_dict is provided, a new
-    dictionary is created.
+    Helper function to load and update macros from the specified directory.
 
     Args:
-        substitutions_dict (dict, optional): A dictionary to be augmented with Jinja macros.
-                                             Defaults to None.
-
-    Returns:
-        dict: The updated substitutions_dict containing the Jinja macros.
+        macros_directory (str): The path to the directory containing macro files.
+        substitutions_dict (dict): A dictionary to be augmented with Jinja macros.
 
     Raises:
         RuntimeError: If there is an error while reading or processing the macro files.
     """
-    if substitutions_dict is None:
-        substitutions_dict = dict()
-
-    add_python_functions(substitutions_dict)
     try:
-        for filename in sorted(os.listdir(JINJA_MACROS_DIRECTORY)):
+        for filename in sorted(os.listdir(macros_directory)):
             if filename.endswith(".jinja"):
-                macros_file = os.path.join(JINJA_MACROS_DIRECTORY, filename)
+                macros_file = os.path.join(macros_directory, filename)
                 update_substitutions_dict(macros_file, substitutions_dict)
     except Exception as exc:
         msg = ("Error extracting macro definitions from '{1}': {0}"
                .format(str(exc), filename))
         raise RuntimeError(msg)
 
+
+def _load_macros(macros_directory, substitutions_dict=None):
+    """
+    Load macros from a specified directory and add them to a substitutions dictionary.
+
+    This function checks if the given macros directory exists, adds Python functions to the
+    substitutions dictionary, and then loads macros from the directory into the dictionary.
+
+    Args:
+        macros_directory (str): The path to the directory containing macro files.
+        substitutions_dict (dict, optional): A dictionary to store the loaded macros.
+                                             If None, a new dictionary is created.
+
+    Returns:
+        dict: The updated substitutions dictionary containing the loaded macros.
+
+    Raises:
+        RuntimeError: If the specified macros directory does not exist.
+    """
+    if substitutions_dict is None:
+        substitutions_dict = dict()
+
+    add_python_functions(substitutions_dict)
+
+    if not os.path.isdir(macros_directory):
+        msg = (f"The directory '{macros_directory}' does not exist.")
+        raise RuntimeError(msg)
+
+    _load_macros_from_directory(macros_directory, substitutions_dict)
+
     return substitutions_dict
 
 
+def load_macros(substitutions_dict=None):
+    """
+    Augments the provided substitutions_dict with project Jinja macros found in the in
+    JINJA_MACROS_DIRECTORY from constants.py.
+
+    Args:
+        substitutions_dict (dict, optional): A dictionary to be augmented with Jinja macros.
+                                             Defaults to None.
+
+    Returns:
+        dict: The updated substitutions_dict containing the Jinja macros.
+    """
+    return _load_macros(JINJA_MACROS_DIRECTORY, substitutions_dict)
+
+
+def load_macros_from_content_dir(content_dir, substitutions_dict=None):
+    """
+    Augments the provided substitutions_dict with project Jinja macros found in a specified
+    content directory.
+
+    Args:
+        content_dir (str): The base directory containing the 'shared/macros' subdirectory.
+        substitutions_dict (dict, optional): A dictionary to be augmented with Jinja macros.
+                                             Defaults to None.
+
+    Returns:
+        dict: The updated substitutions_dict containing the Jinja macros.
+    """
+    jinja_macros_directory = os.path.join(content_dir, 'shared', 'macros')
+    return _load_macros(jinja_macros_directory, substitutions_dict)
+
+
 def process_file_with_macros(filepath, substitutions_dict):
     """
     Process a file with Jinja macros.

ssg/variables.py

@@ -8,7 +8,7 @@
 from collections import defaultdict
 from .constants import BENCHMARKS
 from .profiles import get_profiles_from_products
-from .yaml import open_and_macro_expand
+from .yaml import open_and_macro_expand_from_dir
 
 
 if sys.version_info >= (3, 9):
@@ -82,7 +82,7 @@ def _get_variables_content(content_dir: str) -> dict_type:
 
     for var_file in get_variable_files(content_dir):
         try:
-            yaml_content = open_and_macro_expand(var_file)
+            yaml_content = open_and_macro_expand_from_dir(var_file, content_dir)
         except Exception as e:
             print(f"Error processing file {var_file}: {e}")
             continue

ssg/yaml.py

@@ -12,7 +12,11 @@
 
 from collections import OrderedDict
 
-from .jinja import load_macros, process_file
+from .jinja import (
+    load_macros,
+    load_macros_from_content_dir,
+    process_file,
+)
 
 try:
     from yaml import CSafeLoader as yaml_SafeLoader
@@ -217,6 +221,27 @@ def open_and_macro_expand(yaml_file, substitutions_dict=None):
     return open_and_expand(yaml_file, substitutions_dict)
 
 
+def open_and_macro_expand_from_dir(yaml_file, content_dir, substitutions_dict=None):
+    """
+    Opens a YAML file and expands macros from a specified directory. It is similar to
+    open_and_macro_expand but loads macro definitions from a specified directory instead of the
+    default directory defined in constants. This is useful in cases where the SSG library is
+    consumed by an external project.
+
+    Args:
+        yaml_file (str): The path to the YAML file to be opened and expanded.
+        content_dir (str): The content dir directory to be used for expansion.
+        substitutions_dict (dict, optional): A dictionary of substitutions to be used for macro
+                                             expansion. If None, a new dictionary will be created
+                                             from the content_dir.
+
+    Returns:
+        dict: The expanded content of the YAML file.
+    """
+    substitutions_dict = load_macros_from_content_dir(content_dir, substitutions_dict)
+    return open_and_expand(yaml_file, substitutions_dict)
+
+
 def open_raw(yaml_file):
     """
     Open the given YAML file and parse its contents without performing any template processing.

tests/unit/ssg-module/test_jinja.py

@@ -1,7 +1,5 @@
 import os
-
 import pytest
-
 import ssg.jinja
 
 
@@ -26,3 +24,32 @@ def test_macro_expansion():
 
     complete_defs = get_definitions_with_substitution(dict(global_var="value"))
     assert complete_defs["expand_to_global_var"]() == "value"
+
+
+def test_load_macros_with_valid_directory(tmpdir):
+    macros_dir = tmpdir.mkdir("macros")
+    macro_file = macros_dir.join("test_macro.jinja")
+    macro_file.write("{{% macro test_macro() %}}test{{% endmacro %}}")
+    substitutions_dict = ssg.jinja._load_macros(str(macros_dir))
+
+    assert "test_macro" in substitutions_dict
+    assert substitutions_dict["test_macro"]() == "test"
+
+
+def test_load_macros_with_nonexistent_directory():
+    non_existent_dir = "/non/existent/directory"
+    with pytest.raises(RuntimeError, match=f"The directory '{non_existent_dir}' does not exist."):
+        ssg.jinja._load_macros(non_existent_dir)
+
+
+def test_load_macros_with_existing_substitutions_dict(tmpdir):
+    macros_dir = tmpdir.mkdir("macros")
+    macro_file = macros_dir.join("test_macro.jinja")
+    macro_file.write("{{% macro test_macro() %}}test{{% endmacro %}}")
+    existing_dict = {"existing_key": "existing_value"}
+    substitutions_dict = ssg.jinja._load_macros(str(macros_dir), existing_dict)
+
+    assert "test_macro" in substitutions_dict
+    assert substitutions_dict["test_macro"]() == "test"
+    assert "existing_key" in substitutions_dict
+    assert substitutions_dict["existing_key"] == "existing_value"

tests/unit/ssg-module/test_variables.py

@@ -1,5 +1,4 @@
 import os
-import pytest
 
 from ssg.constants import BENCHMARKS
 from ssg.variables import (
@@ -22,6 +21,10 @@ def setup_test_files(base_dir, benchmark_dirs, create_txt_file=False):
         benchmark_dirs (list[str]): List of benchmark folder paths to create.
         create_txt_file (bool): Whether to create an additional .txt file in each benchmark.
     """
+    # Ensures the shared/macros directory exists even if in this case the testing example does
+    # not use Jinja2 macros.
+    os.makedirs(base_dir / "shared" / "macros", exist_ok=True)
+
     for benchmark_dir in benchmark_dirs:
         path = base_dir / benchmark_dir
         os.makedirs(path, exist_ok=True)
@@ -119,6 +122,7 @@ def __init__(self, product_id, profile_id, variables):
     result = get_variables_from_profiles(profiles)
     assert result == expected_result
 
+
 def test_get_variable_property(tmp_path):
     content_dir = tmp_path / "content"
     benchmark_dirs = ["app", "app/rules"]

tests/unit/ssg-module/test_yaml.py

@@ -1,3 +1,5 @@
+import os
+
 import ssg.yaml
 
 
@@ -46,3 +48,21 @@ def test_list_or_string_update():
         ["something", "entirely"],
         ["entirely", "else"],
     ) == ["something", "entirely", "entirely", "else"]
+
+
+def test_open_and_macro_expand_from_dir(tmpdir):
+    # Setup: Create directory structure
+    content_dir = tmpdir / "content_dir"
+    macros_dir = content_dir / "shared" / "macros"
+    os.makedirs(macros_dir, exist_ok=True)
+
+    # Create YAML file with macro
+    yaml_file = content_dir / "test.yaml"
+    yaml_file.write("macro: {{{ test_macro() }}}")
+
+    # Create macro file with macro definition
+    macro_file = macros_dir / "test_macro.jinja"
+    macro_file.write("{{% macro test_macro() %}}test{{% endmacro %}}")
+
+    result = ssg.yaml.open_and_macro_expand_from_dir(str(yaml_file), str(content_dir))
+    assert result['macro'] == 'test'