add for expressions rules

kkozik-amplify · kkozik-amplify · commit 5932662bfe50 · 2025-09-15T12:26:59.000+02:00
diff --git a/hcl2/rule_transformer/rules/for_expressions.py b/hcl2/rule_transformer/rules/for_expressions.py
@@ -0,0 +1,283 @@
+from typing import Any, Tuple, Optional, List
+
+from lark.tree import Meta
+
+from hcl2.rule_transformer.rules.abstract import LarkRule, LarkElement
+from hcl2.rule_transformer.rules.expressions import ExpressionRule
+from hcl2.rule_transformer.rules.literal_rules import IdentifierRule
+from hcl2.rule_transformer.rules.tokens import (
+    LSQB,
+    RSQB,
+    LBRACE,
+    RBRACE,
+    FOR,
+    IN,
+    IF,
+    COMMA,
+    COLON,
+    ELLIPSIS,
+    FOR_OBJECT_ARROW,
+)
+from hcl2.rule_transformer.rules.whitespace import (
+    NewLineOrCommentRule,
+    InlineCommentMixIn,
+)
+from hcl2.rule_transformer.utils import (
+    SerializationOptions,
+    SerializationContext,
+    to_dollar_string,
+)
+
+
+class ForIntroRule(InlineCommentMixIn):
+    """Rule for the intro part of for expressions: 'for key, value in collection :'"""
+
+    _children: Tuple[
+        FOR,
+        Optional[NewLineOrCommentRule],
+        IdentifierRule,
+        Optional[COMMA],
+        Optional[IdentifierRule],
+        Optional[NewLineOrCommentRule],
+        IN,
+        Optional[NewLineOrCommentRule],
+        ExpressionRule,
+        Optional[NewLineOrCommentRule],
+        COLON,
+        Optional[NewLineOrCommentRule],
+    ]
+
+    @staticmethod
+    def lark_name() -> str:
+        return "for_intro"
+
+    def __init__(self, children, meta: Optional[Meta] = None):
+        # Insert null comments at positions where they might be missing
+        self._possibly_insert_null_second_identifier(children)
+        self._possibly_insert_null_comments(children, [1, 5, 7, 9, 11])
+        super().__init__(children, meta)
+
+    def _possibly_insert_null_second_identifier(self, children: List[LarkRule]):
+        second_identifier_present = (
+            len([child for child in children if isinstance(child, IdentifierRule)]) == 2
+        )
+        if not second_identifier_present:
+            children.insert(3, None)
+            children.insert(4, None)
+
+    @property
+    def first_iterator(self) -> IdentifierRule:
+        """Returns the first iterator"""
+        return self._children[2]
+
+    @property
+    def second_iterator(self) -> Optional[IdentifierRule]:
+        """Returns the second iterator or None if not present"""
+        return self._children[4]
+
+    @property
+    def iterable(self) -> ExpressionRule:
+        """Returns the collection expression being iterated over"""
+        return self._children[8]
+
+    def serialize(
+        self, options=SerializationOptions(), context=SerializationContext()
+    ) -> str:
+        result = "for "
+
+        result += f"{self.first_iterator.serialize(options, context)}"
+        if self.second_iterator:
+            result += f", {self.second_iterator.serialize(options, context)}"
+
+        result += f" in {self.iterable.serialize(options, context)} : "
+
+        return result
+
+
+class ForCondRule(InlineCommentMixIn):
+    """Rule for the optional condition in for expressions: 'if condition'"""
+
+    _children: Tuple[
+        IF,
+        Optional[NewLineOrCommentRule],
+        ExpressionRule,  # condition expression
+    ]
+
+    @staticmethod
+    def lark_name() -> str:
+        return "for_cond"
+
+    def __init__(self, children, meta: Optional[Meta] = None):
+        self._possibly_insert_null_comments(children, [1])
+        super().__init__(children, meta)
+
+    @property
+    def condition_expr(self) -> ExpressionRule:
+        """Returns the condition expression"""
+        return self._children[2]
+
+    def serialize(
+        self, options=SerializationOptions(), context=SerializationContext()
+    ) -> str:
+        return f"if {self.condition_expr.serialize(options, context)}"
+
+
+class ForTupleExprRule(ExpressionRule):
+    """Rule for tuple/array for expressions: [for item in items : expression]"""
+
+    _children: Tuple[
+        LSQB,
+        Optional[NewLineOrCommentRule],
+        ForIntroRule,
+        Optional[NewLineOrCommentRule],
+        ExpressionRule,
+        Optional[NewLineOrCommentRule],
+        Optional[ForCondRule],
+        Optional[NewLineOrCommentRule],
+        RSQB,
+    ]
+
+    @staticmethod
+    def lark_name() -> str:
+        return "for_tuple_expr"
+
+    def __init__(self, children, meta: Optional[Meta] = None):
+        self._possibly_insert_null_comments(children, [1, 3, 5, 7])
+        self._possibly_insert_null_condition(children)
+        super().__init__(children, meta)
+
+    def _possibly_insert_null_condition(self, children: List[LarkElement]):
+        if not len([child for child in children if isinstance(child, ForCondRule)]):
+            children.insert(6, None)
+
+    @property
+    def for_intro(self) -> ForIntroRule:
+        """Returns the for intro rule"""
+        return self._children[2]
+
+    @property
+    def value_expr(self) -> ExpressionRule:
+        """Returns the value expression"""
+        return self._children[4]
+
+    @property
+    def condition(self) -> Optional[ForCondRule]:
+        """Returns the optional condition rule"""
+        return self._children[6]
+
+    def serialize(
+        self, options=SerializationOptions(), context=SerializationContext()
+    ) -> Any:
+
+        result = "["
+
+        with context.modify(inside_dollar_string=True):
+            result += self.for_intro.serialize(options, context)
+            result += self.value_expr.serialize(options, context)
+
+            if self.condition is not None:
+                result += f" {self.condition.serialize(options, context)}"
+
+        result += "]"
+        if not context.inside_dollar_string:
+            result = to_dollar_string(result)
+        return result
+
+
+class ForObjectExprRule(ExpressionRule):
+    """Rule for object for expressions: {for key, value in items : key => value}"""
+
+    _children: Tuple[
+        LBRACE,
+        Optional[NewLineOrCommentRule],
+        ForIntroRule,
+        Optional[NewLineOrCommentRule],
+        ExpressionRule,
+        FOR_OBJECT_ARROW,
+        Optional[NewLineOrCommentRule],
+        ExpressionRule,
+        Optional[NewLineOrCommentRule],
+        Optional[ELLIPSIS],
+        Optional[NewLineOrCommentRule],
+        Optional[ForCondRule],
+        Optional[NewLineOrCommentRule],
+        RBRACE,
+    ]
+
+    @staticmethod
+    def lark_name() -> str:
+        return "for_object_expr"
+
+    def __init__(self, children, meta: Optional[Meta] = None):
+        self._possibly_insert_null_comments(children, [1, 3, 6, 8, 10, 12])
+        self._possibly_insert_null_optionals(children)
+        super().__init__(children, meta)
+
+    def _possibly_insert_null_optionals(self, children: List[LarkElement]):
+        has_ellipsis = False
+        has_condition = False
+
+        for child in children:
+            # if not has_ellipsis and isinstance(child, ELLIPSIS):
+            if (
+                has_ellipsis is False
+                and child is not None
+                and child.lark_name() == ELLIPSIS.lark_name()
+            ):
+                has_ellipsis = True
+            if not has_condition and isinstance(child, ForCondRule):
+                has_condition = True
+
+        if not has_ellipsis:
+            children.insert(9, None)
+
+        if not has_condition:
+            children.insert(11, None)
+
+    @property
+    def for_intro(self) -> ForIntroRule:
+        """Returns the for intro rule"""
+        return self._children[2]
+
+    @property
+    def key_expr(self) -> ExpressionRule:
+        """Returns the key expression"""
+        return self._children[4]
+
+    @property
+    def value_expr(self) -> ExpressionRule:
+        """Returns the value expression"""
+        return self._children[7]
+
+    @property
+    def ellipsis(self) -> Optional[ELLIPSIS]:
+        """Returns the optional ellipsis token"""
+        return self._children[9]
+
+    @property
+    def condition(self) -> Optional[ForCondRule]:
+        """Returns the optional condition rule"""
+        return self._children[11]
+
+    def serialize(
+        self, options=SerializationOptions(), context=SerializationContext()
+    ) -> Any:
+        result = "{"
+        with context.modify(inside_dollar_string=True):
+            result += self.for_intro.serialize(options, context)
+            result += f"{self.key_expr.serialize(options, context)} => "
+
+            result += self.value_expr.serialize(
+                SerializationOptions(wrap_objects=True), context
+            )
+
+            if self.ellipsis is not None:
+                result += self.ellipsis.serialize(options, context)
+
+            if self.condition is not None:
+                result += f" {self.condition.serialize(options, context)}"
+
+        result += "}"
+        if not context.inside_dollar_string:
+            result = to_dollar_string(result)
+        return result
diff --git a/hcl2/rule_transformer/rules/functions.py b/hcl2/rule_transformer/rules/functions.py
@@ -40,7 +40,7 @@ def arguments(self) -> List[ExpressionRule]:
         return [child for child in self._children if isinstance(child, ExpressionRule)]
 
     def serialize(self, options = SerializationOptions(), context = SerializationContext()) -> Any:
-        result = ", ".join([argument.serialize(options, context) for argument in self.arguments])
+        result = ", ".join([str(argument.serialize(options, context)) for argument in self.arguments])
         if self.has_ellipsis:
             result += " ..."
         return result
diff --git a/hcl2/rule_transformer/rules/tokens.py b/hcl2/rule_transformer/rules/tokens.py
@@ -91,6 +91,10 @@ def serialize_conversion(self) -> Callable[[Any], str]:
 DBLQUOTE = StaticStringToken[("DBLQUOTE", '"')]
 ATTR_SPLAT = StaticStringToken[("ATTR_SPLAT", ".*")]
 FULL_SPLAT = StaticStringToken[("FULL_SPLAT", "[*]")]
+FOR = StaticStringToken[("FOR", "for")]
+IN = StaticStringToken[("IN", "in")]
+IF = StaticStringToken[("IF", "if")]
+FOR_OBJECT_ARROW = StaticStringToken[("FOR_OBJECT_ARROW", "=>")]
 
 
 class IntLiteral(LarkToken):
diff --git a/hcl2/rule_transformer/transformer.py b/hcl2/rule_transformer/transformer.py
@@ -23,6 +23,12 @@
     ExprTermRule,
     ConditionalRule,
 )
+from hcl2.rule_transformer.rules.for_expressions import (
+    ForTupleExprRule,
+    ForObjectExprRule,
+    ForIntroRule,
+    ForCondRule,
+)
 from hcl2.rule_transformer.rules.functions import ArgumentsRule, FunctionCallRule
 from hcl2.rule_transformer.rules.indexing import (
     IndexExprTermRule,
@@ -40,12 +46,13 @@
     IntLitRule,
     IdentifierRule,
     BinaryOperatorRule,
+    KeywordRule,
 )
 from hcl2.rule_transformer.rules.strings import (
     InterpolationRule,
     StringRule,
-    StringPartRule, 
-    HeredocTemplateRule, 
+    StringPartRule,
+    HeredocTemplateRule,
     HeredocTrimTemplateRule,
 )
 from hcl2.rule_transformer.rules.tokens import (
@@ -72,6 +79,7 @@ def __init__(self, discard_new_line_or_comments: bool = False):
         self.discard_new_line_or_comments = discard_new_line_or_comments
 
     def __default_token__(self, token: Token) -> StringToken:
+        # TODO make this return StaticStringToken where applicable
         return StringToken[token.type](token.value)
 
     def FLOAT_LITERAL(self, token: Token) -> FloatLiteral:
@@ -109,6 +117,10 @@ def new_line_or_comment(self, meta: Meta, args) -> NewLineOrCommentRule:
     def identifier(self, meta: Meta, args) -> IdentifierRule:
         return IdentifierRule(args, meta)
 
+    @v_args(meta=True)
+    def keyword(self, meta: Meta, args) -> KeywordRule:
+        return KeywordRule(args, meta)
+
     @v_args(meta=True)
     def int_lit(self, meta: Meta, args) -> IntLitRule:
         return IntLitRule(args, meta)
@@ -132,11 +144,11 @@ def interpolation(self, meta: Meta, args) -> InterpolationRule:
     @v_args(meta=True)
     def heredoc_template(self, meta: Meta, args) -> HeredocTemplateRule:
         return HeredocTemplateRule(args, meta)
-        
+
     @v_args(meta=True)
     def heredoc_template_trim(self, meta: Meta, args) -> HeredocTrimTemplateRule:
         return HeredocTrimTemplateRule(args, meta)
-    
+
     @v_args(meta=True)
     def expr_term(self, meta: Meta, args) -> ExprTermRule:
         return ExprTermRule(args, meta)
@@ -236,3 +248,19 @@ def full_splat(self, meta: Meta, args) -> FullSplatRule:
     @v_args(meta=True)
     def full_splat_expr_term(self, meta: Meta, args) -> FullSplatExprTermRule:
         return FullSplatExprTermRule(args, meta)
+
+    @v_args(meta=True)
+    def for_tuple_expr(self, meta: Meta, args) -> ForTupleExprRule:
+        return ForTupleExprRule(args, meta)
+
+    @v_args(meta=True)
+    def for_object_expr(self, meta: Meta, args) -> ForObjectExprRule:
+        return ForObjectExprRule(args, meta)
+
+    @v_args(meta=True)
+    def for_intro(self, meta: Meta, args) -> ForIntroRule:
+        return ForIntroRule(args, meta)
+
+    @v_args(meta=True)
+    def for_cond(self, meta: Meta, args) -> ForCondRule:
+        return ForCondRule(args, meta)
