Refactor spec parser and move utility to template engine

Author: polyaxon-ci

cli/polyaxon/_polyaxonfile/specs/libs/engine.py

@@ -1,9 +1,11 @@
 import jinja2
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from markupsafe import soft_str
 
+from polyaxon.exceptions import PolyaxonSchemaError
+
 
 def map_format(value, pattern, variable_name: Optional[str] = None):
     if variable_name:
@@ -15,3 +17,69 @@ def get_engine():
     env = jinja2.Environment()
     env.filters["map_format"] = map_format
     return env
+
+
+def render_template(template_str: str, context: Dict[str, Any]) -> str:
+    """
+    Render a Jinja2 template string with context.
+
+    Args:
+        template_str: Template string with Jinja2 syntax
+        context: Dictionary of variables for template rendering
+
+    Returns:
+        Rendered string
+
+    Raises:
+        PolyaxonSchemaError: If template rendering fails
+    """
+    if not template_str or not isinstance(template_str, str):
+        return template_str
+
+    if "{{" not in template_str and "{%" not in template_str:
+        return template_str  # No template syntax, return as-is
+
+    try:
+        engine = get_engine()
+        template = engine.from_string(template_str)
+        return template.render(**context)
+    except (jinja2.exceptions.TemplateError, ValueError, TypeError) as e:
+        raise PolyaxonSchemaError(
+            f"Encountered a problem rendering the template, "
+            f"please make sure your variables are resolvable. "
+            f"Error: {repr(e)}"
+        )
+
+
+def render_config(config: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Recursively render all string values in a config dict.
+
+    Args:
+        config: Configuration dictionary with potential template strings
+        context: Dictionary of variables for template rendering
+
+    Returns:
+        Config with all template strings rendered
+    """
+    if not config:
+        return config
+
+    rendered = {}
+    for key, value in config.items():
+        if isinstance(value, str):
+            rendered[key] = render_template(value, context)
+        elif isinstance(value, dict):
+            rendered[key] = render_config(value, context)
+        elif isinstance(value, list):
+            rendered[key] = [
+                render_config(item, context)
+                if isinstance(item, dict)
+                else render_template(item, context)
+                if isinstance(item, str)
+                else item
+                for item in value
+            ]
+        else:
+            rendered[key] = value
+    return rendered

cli/polyaxon/_polyaxonfile/specs/libs/parser.py

@@ -1,5 +1,4 @@
 import ast
-import jinja2
 
 from collections.abc import Mapping
 from datetime import date, datetime, timedelta
@@ -15,9 +14,8 @@
 )
 
 from polyaxon._flow import ParamSpec
-from polyaxon._polyaxonfile.specs.libs.engine import get_engine
+from polyaxon._polyaxonfile.specs.libs.engine import render_template
 from polyaxon._polyaxonfile.specs.sections import Sections
-from polyaxon.exceptions import PolyaxonSchemaError
 
 try:
     import numpy as np
@@ -28,8 +26,6 @@
 class PolyaxonfileParser:
     """Parses the Polyaxonfile."""
 
-    engine = get_engine()
-
     @staticmethod
     def _get_section_data(section_data):
         if hasattr(section_data, "to_dict"):
@@ -193,19 +189,6 @@ def parse_io(
     @classmethod
     def parse_expression(  # pylint:disable=too-many-branches
         cls, expression, params: Dict, check_operators: bool = False
-    ):
-        try:
-            return cls._parse_expression(expression, params, check_operators)
-        except jinja2.exceptions.TemplateError as e:
-            raise PolyaxonSchemaError(
-                "Encountered a problem parsing the template, "
-                "please make sure your variables are resolvable. "
-                "Error: {}".format(repr(e))
-            )
-
-    @classmethod
-    def _parse_expression(  # pylint:disable=too-many-branches
-        cls, expression, params: Dict, check_operators: bool = False
     ):
         if isinstance(expression, (int, float, complex, type(None))):
             return expression
@@ -227,7 +210,7 @@ def _parse_expression(  # pylint:disable=too-many-branches
             if len(expression) == 1:
                 old_key, value = list(expression.items())[0]
                 # always parse the keys, they must be base object or evaluate to base objects
-                key = cls._parse_expression(old_key, params)
+                key = cls.parse_expression(old_key, params)
                 if check_operators and cls.is_operator(key):
                     return cls._parse_operator({key: value}, params)
                 else:
@@ -236,7 +219,7 @@ def _parse_expression(  # pylint:disable=too-many-branches
             new_expression = {}
             for k, v in expression.items():
                 new_expression.update(
-                    cls._parse_expression({k: v}, params, check_operators)
+                    cls.parse_expression({k: v}, params, check_operators)
                 )
             return new_expression
 
@@ -253,14 +236,7 @@ def _parse_expression(  # pylint:disable=too-many-branches
 
     @classmethod
     def _evaluate_expression(cls, expression, params, check_operators):
-        try:
-            result = cls.engine.from_string(expression).render(**params)
-        except (ValueError, TypeError) as e:
-            raise PolyaxonSchemaError(
-                "Encountered a problem parsing the template, "
-                "please make sure your variables are resolvable. "
-                "Error: {}".format(repr(e))
-            )
+        result = render_template(expression, params)
         if result == expression:
             try:
                 return ast.literal_eval(result)