From c291a3f5cb1bad803fa1055f2a3fcf251d37a2a4 Mon Sep 17 00:00:00 2001
From: Akemi Izuko <akemi@noway.moe>
Date: Wed, 17 Jul 2024 19:24:52 -0600
Subject: [PATCH] Parser: update for loop syntax

---
 src/main.py | 96 +++++++++++++++++------------------------------------
 1 file changed, 31 insertions(+), 65 deletions(-)

diff --git a/src/main.py b/src/main.py
index b1a4496..7887483 100644
--- a/src/main.py
+++ b/src/main.py
@@ -42,37 +42,32 @@ def run_assertions(host, props):
 # │ Pαrsεr                                                                    |
 # ╚───────────────────────────────────────────────────────────────────────────╝
 def expand_for_loop(props):
-    expand = list()
+    expanded_props = dict()
 
     for loop in props.get("for", []):
-        r = loop["range"]
+        prop = loop["property"]
 
-        if isinstance(r, dict):
-            start = int(r["from"])
-            cease = int(r["to"])
-            inc = int(r.get("increment", 1))
+        if expanded_props.get(prop) is None:
+            expanded_props[prop] = list()
 
-            iterations = range(start, cease, inc)
-        elif isinstance(r, list):
-            iterations = r
+        if loop.get("range"):
+            r = loop["range"]
+            start = int(r[0])
+            cease = int(r[1])
+            incre = int(r[2]) if len(r) > 2 else 1
+
+            iterations = range(start, cease, incre)
+        elif loop.get("iter"):
+            iterations = loop["iter"]
         else:
-            raise TypeError("Expected list or dict in `for` loop `range`")
+            raise TypeError("Expected a `range` or `iter` property on `for`")
 
         pattern = re.compile(rf'\$\{{{loop["variable"]}\}}')
 
         for i in iterations:
-            expand.append(pattern.sub(str(i), loop["template"]))
+            expanded_props[prop].append(pattern.sub(str(i), loop["template"]))
 
-    return expand
-
-
-def get_ssh_props(data):
-    props = dict()
-
-    for k, v in data["ssh_props"].items():
-        props[k] = v if isinstance(v, list) else [v]
-
-    return props
+    return expanded_props
 
 
 def parse_props(host, props, hosts):
@@ -83,25 +78,11 @@ def parse_props(host, props, hosts):
         parsed[k] = v if isinstance(v, list) else [v]
 
     # For-loops in config
-    for loop in props.get("for", []):
-        prop = loop["property"]
-        r = loop["range"]
-
-        if isinstance(r, dict):
-            iterations = range(r["from"], r["to"])
-        elif isinstance(r, list):
-            iterations = r
+    for prop, values in expand_for_loop(props).items():
+        if parsed.get(prop) is not None:
+            parsed[prop].extend(values)
         else:
-            raise TypeError("Expected list or dict in `for` loop `range`")
-
-        if not parsed.get(prop):
-            parsed[prop] = list()
-
-        pattern = re.compile(rf'\$\{{{loop["variable"]}\}}')
-
-        for i in iterations:
-            sub = pattern.sub(str(i), loop["template"])
-            parsed[prop].append(sub)
+            parsed[prop] = values
 
     # Inherit from template
     template = props.get("template")
@@ -128,22 +109,6 @@ def remove_templates(expanded):
     return hosts
 
 
-def stringify_booleans(props):
-    pattern = re.compile(r"(True|False)$")
-
-    def replace_bools(s):
-        m = pattern.search(s)
-
-        if m and m[1] == "True":
-            return pattern.sub("true", s)
-        elif m and m[1] == "False":
-            return pattern.sub("false", s)
-        else:
-            return s
-
-    return [replace_bools(prop) for prop in props]
-
-
 def parse_yaml(data):
     hosts = dict()
     check_duplicate_hosts(data)
@@ -180,15 +145,16 @@ def to_ssh_config_string(parse):
 # ╔───────────────────────────────────────────────────────────────────────────╗
 # │ Mαiη                                                                      |
 # ╚───────────────────────────────────────────────────────────────────────────╝
-with open(args.file, "r") as f:
-    data = yaml.safe_load(f)
+if __name__ == '__main__':
+    with open(args.file, "r") as f:
+        data = yaml.safe_load(f)
 
-parse = parse_yaml(data)
+    parse = parse_yaml(data)
 
-if args.json:
-    print(json.dumps(parse))
-elif args.out:
-    with open(args.out, "w") as f:
-        f.write(to_ssh_config_string(parse))
-else:
-    print(to_ssh_config_string(parse), end="")
+    if args.json:
+        print(json.dumps(parse))
+    elif args.out:
+        with open(args.out, "w") as f:
+            f.write(to_ssh_config_string(parse))
+    else:
+        print(to_ssh_config_string(parse), end="")