add Lark AST -> HCL2 reconstructor and LarkTree formatter; various other fixes and changes:

* preserve order of serialized attributes and blocks
* make RuleTransformer.__default_token__ differentiate between StringToken and StaticStringToken
* add separate ProviderFunctionCallRule class for more accurate reconstruction

Author: kkozik-amplify

hcl2/rule_transformer/deserializer.py

@@ -1,6 +1,8 @@
 import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from functools import lru_cache
-from typing import Any, TextIO, List, Union
+from typing import Any, TextIO, List, Union, Optional
 
 from regex import regex
 
@@ -31,7 +33,7 @@
     StringRule,
     InterpolationRule,
     StringPartRule,
-    HeredocTemplateRule, 
+    HeredocTemplateRule,
     HeredocTrimTemplateRule,
 )
 from hcl2.rule_transformer.rules.tokens import (
@@ -51,34 +53,62 @@
     LBRACE,
     HEREDOC_TRIM_TEMPLATE,
     HEREDOC_TEMPLATE,
+    COLON,
 )
+from hcl2.rule_transformer.rules.whitespace import NewLineOrCommentRule
 from hcl2.rule_transformer.transformer import RuleTransformer
-from hcl2.rule_transformer.utils import DeserializationOptions, HEREDOC_TRIM_PATTERN, HEREDOC_PATTERN
+from hcl2.rule_transformer.utils import HEREDOC_TRIM_PATTERN, HEREDOC_PATTERN
 
 
-class Deserializer:
-    def __init__(self, options=DeserializationOptions()):
-        self.options = options
+@dataclass
+class DeserializerOptions:
+    heredocs_to_strings: bool = False
+    indent_length: int = 2
+    object_elements_colon: bool = False
+    object_elements_trailing_comma: bool = True
+
+
+class LarkElementTreeDeserializer(ABC):
+    def __init__(self, options: DeserializerOptions = None):
+        self.options = options or DeserializerOptions()
+
+    @abstractmethod
+    def loads(self, value: str) -> LarkElement:
+        raise NotImplementedError()
+
+    def load(self, file: TextIO) -> LarkElement:
+        return self.loads(file.read())
+
+
+class BaseDeserializer(LarkElementTreeDeserializer):
+    def __init__(self, options=None):
+        super().__init__(options)
+        self._current_line = 1
+        self._last_new_line: Optional[NewLineOrCommentRule] = None
 
     @property
     @lru_cache
     def _transformer(self) -> RuleTransformer:
         return RuleTransformer()
 
     def load_python(self, value: Any) -> LarkElement:
-        return StartRule([self._deserialize(value)])
+        result = StartRule([self._deserialize(value)])
+        return result
 
     def loads(self, value: str) -> LarkElement:
         return self.load_python(json.loads(value))
 
-    def load(self, file: TextIO) -> LarkElement:
-        return self.loads(file.read())
-
     def _deserialize(self, value: Any) -> LarkElement:
         if isinstance(value, dict):
             if self._contains_block_marker(value):
-                elements = self._deserialize_block_elements(value)
-                return BodyRule(elements)
+
+                children = []
+
+                block_elements = self._deserialize_block_elements(value)
+                for element in block_elements:
+                    children.append(element)
+
+                return BodyRule(children)
 
             return self._deserialize_object(value)
 
@@ -89,14 +119,13 @@ def _deserialize(self, value: Any) -> LarkElement:
 
     def _deserialize_block_elements(self, value: dict) -> List[LarkRule]:
         children = []
-
         for key, value in value.items():
             if self._is_block(value):
                 # this value is a list of blocks, iterate over each block and deserialize them
                 for block in value:
                     children.append(self._deserialize_block(key, block))
-            else:
 
+            else:
                 # otherwise it's just an attribute
                 if key != IS_BLOCK:
                     children.append(self._deserialize_attribute(key, value))
@@ -106,28 +135,24 @@ def _deserialize_block_elements(self, value: dict) -> List[LarkRule]:
     def _deserialize_text(self, value: Any) -> LarkRule:
         try:
             int_val = int(value)
+            if "." in str(value):
+                return FloatLitRule([FloatLiteral(float(value))])
             return IntLitRule([IntLiteral(int_val)])
         except ValueError:
             pass
 
-        try:
-            float_val = float(value)
-            return FloatLitRule([FloatLiteral(float_val)])
-        except ValueError:
-            pass
-
         if isinstance(value, str):
             if value.startswith('"') and value.endswith('"'):
                 if not self.options.heredocs_to_strings and value.startswith('"<<-'):
                     match = HEREDOC_TRIM_PATTERN.match(value[1:-1])
                     if match:
                         return self._deserialize_heredoc(value[1:-1], True)
-                    
+
                 if not self.options.heredocs_to_strings and value.startswith('"<<'):
                     match = HEREDOC_PATTERN.match(value[1:-1])
                     if match:
                         return self._deserialize_heredoc(value[1:-1], False)
-                    
+
                 return self._deserialize_string(value)
 
             if self._is_expression(value):
@@ -151,7 +176,6 @@ def _deserialize_string(self, value: str) -> StringRule:
         pattern = regex.compile(r"(\${1,2}\{(?:[^{}]|(?R))*\})")
         parts = [part for part in pattern.split(value) if part != ""]
 
-
         for part in parts:
             if part == '"':
                 continue
@@ -181,10 +205,12 @@ def _deserialize_string_part(self, value: str) -> StringPartRule:
 
         return StringPartRule([STRING_CHARS(value)])
 
-    def _deserialize_heredoc(self, value: str, trim: bool) -> Union[HeredocTemplateRule, HeredocTrimTemplateRule]:
+    def _deserialize_heredoc(
+        self, value: str, trim: bool
+    ) -> Union[HeredocTemplateRule, HeredocTrimTemplateRule]:
         if trim:
             return HeredocTrimTemplateRule([HEREDOC_TRIM_TEMPLATE(value)])
-        return HeredocTemplateRule([HEREDOC_TEMPLATE(value)]) 
+        return HeredocTemplateRule([HEREDOC_TEMPLATE(value)])
 
     def _deserialize_expression(self, value: str) -> ExprTermRule:
         """Deserialize an expression string into an ExprTermRule."""
@@ -200,7 +226,9 @@ def _deserialize_expression(self, value: str) -> ExprTermRule:
         # transform parsed tree into LarkElement tree
         rules_tree = self._transformer.transform(parsed_tree)
         # extract expression from the tree
-        return rules_tree.body.children[0].expression
+        result = rules_tree.body.children[0].expression
+
+        return result
 
     def _deserialize_block(self, first_label: str, value: dict) -> BlockRule:
         """Deserialize a block by extracting labels and body"""
@@ -220,14 +248,24 @@ def _deserialize_block(self, first_label: str, value: dict) -> BlockRule:
                 break
 
         return BlockRule(
-            [*[self._deserialize(label) for label in labels], self._deserialize(body)]
+            [
+                *[self._deserialize(label) for label in labels],
+                LBRACE(),
+                self._deserialize(body),
+                RBRACE(),
+            ]
         )
 
     def _deserialize_attribute(self, name: str, value: Any) -> AttributeRule:
+        expr_term = self._deserialize(value)
+
+        if not isinstance(expr_term, ExprTermRule):
+            expr_term = ExprTermRule([expr_term])
+
         children = [
             self._deserialize_identifier(name),
             EQ(),
-            ExprTermRule([self._deserialize(value)]),
+            expr_term,
         ]
         return AttributeRule(children)
 
@@ -247,11 +285,18 @@ def _deserialize_object(self, value: dict) -> ObjectRule:
         children = []
         for key, value in value.items():
             children.append(self._deserialize_object_elem(key, value))
+
         return ObjectRule([LBRACE(), *children, RBRACE()])
 
     def _deserialize_object_elem(self, key: str, value: Any) -> ObjectElemRule:
         if self._is_expression(key):
-            key = ObjectElemKeyExpressionRule([self._deserialize_expression(key)])
+            key = ObjectElemKeyExpressionRule(
+                [
+                    child
+                    for child in self._deserialize_expression(key).children
+                    if child is not None
+                ]
+            )
         elif "." in key:
             parts = key.split(".")
             children = []
@@ -262,13 +307,14 @@ def _deserialize_object_elem(self, key: str, value: Any) -> ObjectElemRule:
         else:
             key = self._deserialize_text(key)
 
-        return ObjectElemRule(
-            [
-                ObjectElemKeyRule([key]),
-                EQ(),
-                ExprTermRule([self._deserialize_text(value)]),
-            ]
-        )
+        result = [
+            ObjectElemKeyRule([key]),
+            COLON() if self.options.object_elements_colon else EQ(),
+            ExprTermRule([self._deserialize(value)]),
+        ]
+        if self.options.object_elements_trailing_comma:
+            result.append(COMMA())
+        return ObjectElemRule(result)
 
     def _is_expression(self, value: str) -> bool:
         return value.startswith("${") and value.endswith("}")