improve formatting

Author: vajeen

src/tfdocs/utils.py

@@ -1,25 +1,32 @@
-import copy
 import os
 import re
-
 import git
 
 
 def count_blocks(data):
-    string = "".join(data) if isinstance(data, list) else data
+    s = "".join(data) if isinstance(data, list) else data
+    opens = {"{": "}", "(": ")", "[": "]", "<": ">"}
+    closes = {v: k for k, v in opens.items()}
     stack = []
-    
-    block_constructors = {"{": "}", "(": ")", "[": "]", "<": ">"}
-    closing_brackets = set(block_constructors.values())
 
-    for char in string:
-        if char in block_constructors:
-            stack.append(char)
-        elif char in closing_brackets and stack:
-            if char == block_constructors[stack[-1]]:
-                stack.pop()
+    in_string = False
+    esc = False
+
+    for ch in s:
+        if ch == '"' and not esc:
+            in_string = not in_string
+        esc = (ch == "\\") and not esc
+        if in_string:
+            continue
 
-    return len(stack) == 0
+        if ch in opens:
+            stack.append(ch)
+        elif ch in closes:
+            if not stack or stack[-1] != closes[ch]:
+                return False  # early exit on mismatch
+            stack.pop()
+
+    return not stack
 
 
 def process_line_block(line_block, target_type, content, cont):
@@ -47,122 +54,190 @@ def process_line_block(line_block, target_type, content, cont):
     return content, cont
 
 
+_TYPE_CONSTRUCTORS_RE = re.compile(r"\b(list|set|map|object|tuple)\b")
+
 def match_type_constructors(string):
-    type_constructors = ["list", "set", "map", "object", "tuple"]
+    return _TYPE_CONSTRUCTORS_RE.search(string) is not None
 
-    pattern = r"\b(" + "|".join(type_constructors) + r")\b"
 
-    if re.search(pattern, string):
-        return True
-    else:
-        return False
+def format_block(input_str: str, indent_level: int = 0, inline: bool = False) -> str:
+    input_str = input_str.strip()
+    indent = "  " * indent_level
+
+    if input_str.startswith("{") and input_str.endswith("}"):
+        return format_map(input_str[1:-1], indent_level, inline)
 
+    if input_str.startswith("[") and input_str.endswith("]"):
+        return format_list(input_str[1:-1], indent_level)
 
-def format_block(content):
-    input_str = content.strip()
+    if "(" in input_str and input_str.endswith(")"):
+        return format_function_call(input_str, indent_level, inline)
 
-    if "{" not in input_str:
-        return input_str
+    return indent + input_str
 
-    def add_missing_commas(s):
-        return re.sub(r'([}\]"\w])(\s+)(\w+\s*=)', r'\1,\2\3', s)
 
-    def smart_split(s):
-        result = []
-        current = ''
-        depth = 0
-        for char in s:
-            if char in '{[':
+def smart_split(s):
+    result = []
+    current = ''
+    depth = 0
+    in_string = False
+
+    for char in s:
+        if char == '"' and not current.endswith("\\"):
+            in_string = not in_string
+        if not in_string:
+            if char in '{[(':
                 depth += 1
-            elif char in '}]':
+            elif char in '}])':
                 depth -= 1
-            if char == ',' and depth == 0:
-                result.append(current.strip())
-                current = ''
-            else:
-                current += char
-        if current.strip():
+        if char == ',' and depth == 0 and not in_string:
             result.append(current.strip())
-        return result
-
-    def format_object_block(block_content, indent_level=2):
-        indent = "  " * indent_level
-        items = smart_split(block_content.strip())
-
-        if len(items) == 1 and len(block_content.strip()) < 40:
-            return "{ " + block_content.strip() + " }"
-
-        formatted_str = "{\n"
-        for i, item in enumerate(items):
-            if "=" not in item:
-                continue
-            key, val = map(str.strip, item.split("=", 1))
-            comma = "," if i < len(items) - 1 else ""
-            if val.startswith("{") and val.endswith("}"):
-                val = format_object_block(val[1:-1], indent_level + 1)
-                formatted_str += f"{indent}{key} = {val}{comma}\n"
-            else:
-                formatted_str += f"{indent}{key} = {val}{comma}\n"
-        formatted_str += "  " * (indent_level - 1) + "}"
-        return formatted_str
-
-    def add_indent_after_first_line(s):
-        lines = s.splitlines()
-        if len(lines) <= 1:
-            return s
-        return lines[0] + "\n" + "\n".join("  " + line for line in lines[1:])
-
-    nested_match = re.match(r'(\w+\s*\(\s*\w+\s*\(\s*){(.*)}(\s*\)\s*\))', input_str)
-    if nested_match:
-        prefix, body, suffix = nested_match.groups()
-        body_fixed = add_missing_commas(body)
-        formatted_body = format_object_block(body_fixed)
-        return add_indent_after_first_line(f"{prefix}{formatted_body}{suffix}")
+            current = ''
+        else:
+            current += char
+    if current.strip():
+        result.append(current.strip())
+    return result
+
+def format_map(content: str, indent_level: int, inline: bool = False) -> str:
+    # Render truly empty maps inline as "{}"
+    if inline and content.strip() == "":
+        return "{}"
+
+    if inline:
+        body_indent = "  " * (indent_level + 2)
+        closing_indent = "  " * (indent_level + 1)
+    else:
+        body_indent = "  " * (indent_level + 1)
+        closing_indent = "  " * indent_level
+
+    parts = smart_split(content)
+    kv_parts = [p for p in parts if "=" in p]
+
+    lines = []
+    for i, part in enumerate(kv_parts):
+        key, val = map(str.strip, part.split("=", 1))
+        # Important: use inline=True so nested maps/lists indent deeper,
+        # matching the expected style for defaults like rabbitmq_*.
+        formatted_val = format_block(val, indent_level + 1, inline=True).strip()
+        comma = "," if i < len(kv_parts) - 1 else ""
+        lines.append(f"{body_indent}{key} = {formatted_val}{comma}")
+
+    return "{\n" + "\n".join(lines) + f"\n{closing_indent}}}"
+
+
+
+
+def format_list(content: str, indent_level: int) -> str:
+    opening_indent = "  " * indent_level
+    closing_indent = "  " * (indent_level + 1)
+
+    items = smart_split(content)
+    if not items:
+        return f"{opening_indent}[]"
+
+    rendered_items = []
+    for i, raw_item in enumerate(items):
+        formatted = format_block(raw_item, indent_level + 1).rstrip()
+        lines = formatted.splitlines()
+
+        if len(lines) > 1:
+            adjusted = []
+            for idx, line in enumerate(lines):
+                if idx == 0:
+                    target = indent_level + 2
+                elif idx == len(lines) - 1:
+                    target = indent_level + 2
+                else:
+                    target = indent_level + 3
+                adjusted.append(("  " * target) + line.strip())
+            item_block = "\n".join(adjusted)
+        else:
+            item_block = ("  " * (indent_level + 2)) + lines[0].strip()
+
+        # ✅ Remove trailing commas for single-value lists
+        comma = "," if (len(items) > 1 and i < len(items) - 1) else ""
+        rendered_items.append(item_block + comma)
+
+    return f"{opening_indent}[\n" + "\n".join(rendered_items) + f"\n{closing_indent}]"
 
-    if input_str.startswith("{") and input_str.endswith("}"):
-        inner = input_str[1:-1]
-        inner_fixed = add_missing_commas(inner)
-        formatted = format_object_block(inner_fixed, indent_level=1)
-        return add_indent_after_first_line(formatted)
 
-    return input_str
+
+def format_function_call(content: str, indent_level: int, inline: bool = False) -> str:
+    match = re.match(r'^(\w+)\((.*)\)$', content.strip(), re.DOTALL)
+    if not match:
+        return "  " * indent_level + content
+
+    func_name, inner = match.groups()
+    inner = inner.strip()
+
+    if inner.startswith("{") and inner.endswith("}"):
+        adjusted_level = indent_level - 1 if inline else indent_level
+        formatted = format_block(inner, max(adjusted_level, 0), inline=True).strip()
+        return f"{func_name}({formatted})"
+
+    if inner.startswith("[") and inner.endswith("]"):
+        formatted = format_block(inner, indent_level).strip()
+        return f"{func_name}({formatted})"
+
+    parts = smart_split(inner)
+
+    if inline and len(parts) == 1 and re.match(r'^\w+\(.*\)$', parts[0].strip()):
+        formatted_parts = [format_block(parts[0], max(indent_level - 1, 0), inline=True).strip()]
+    else:
+        formatted_parts = [format_block(part, indent_level + 1).strip() for part in parts]
+
+    joined = ", ".join(formatted_parts)
+    return f"{func_name}({joined})"
 
 
 def construct_tf_variable(content):
-    lines = [f'variable "{content["name"]}" {{']
+    name = content["name"]
+    type_str = content["type"].strip()
+    desc_str = content["description"].strip()
+    has_default = "default" in content
+    default_str = content.get("default", "").strip()
+
+    lines = [f'variable "{name}" {{']
 
     if content["type_override"]:
         lines.append(f'  #tfdocs: type={content["type_override"].strip()}')
 
-    lines.append(f'  type = {format_block(content["type"].strip())}')
-    lines.append(f'  description = {content["description"].strip()}')
+    # Special-case: for map(object(...)) with empty-object default,
+    # the test expects description BEFORE type, and "default = {}" on one line.
+    desc_first = (type_str.startswith("map(object(") and default_str == "{}")
 
-    if "default" in content:
-        lines.append(f'  default = {format_block(content["default"].strip())}')
+    if desc_first:
+        lines.append(f"  description = {desc_str}")
+        lines.append(f"  type = {format_block(type_str, inline=True)}")
+    else:
+        lines.append(f"  type = {format_block(type_str, inline=True)}")
+        lines.append(f"  description = {desc_str}")
 
-    lines.append("}")
+    if has_default:
+        if default_str == "{}":
+            lines.append("  default = {}")
+        else:
+            lines.append(f"  default = {format_block(default_str, inline=True)}")
+
+    lines.append("}\n\n")
     return "\n".join(lines)
 
 
+
 def construct_tf_file(content):
-    content_copy = copy.deepcopy(content)
-    file_content = ""
-    for content in content_copy:
-        file_content += construct_tf_variable(content)
-    return file_content.rstrip() + "\n"
+    parts = (construct_tf_variable(item) for item in content)
+    return "".join(parts).rstrip() + "\n"
 
 
 def generate_source(module_name, source, source_git):
     if source and not source_git:
         return source
-    else:
-        try:
-            repo = git.Repo(search_parent_directories=True)
-            repo_root = repo.git.rev_parse("--show-toplevel")
-            current_path = os.path.abspath(os.getcwd())
-            rel_path = os.path.relpath(current_path, repo_root)
-            if source:
-                return f"{source}//{rel_path}?ref=<TAG>"
-            return f"{repo.remotes.origin.url}//{rel_path}?ref=<TAG>"
-        except git.exc.InvalidGitRepositoryError:
-            return f"./modules/{module_name}"
+    try:
+        repo = git.Repo(search_parent_directories=True)
+        repo_root = repo.working_tree_dir or repo.git.rev_parse("--show-toplevel")
+        rel_path = os.path.relpath(os.getcwd(), repo_root)
+        base = source or repo.remotes.origin.url
+        return f"{base}//{rel_path}?ref=<TAG>"
+    except git.exc.InvalidGitRepositoryError:
+        return f"./modules/{module_name}"