From 9d11d2ca71e8fcf3c56b5e17e5ff73650bb63f6d Mon Sep 17 00:00:00 2001
From: Faerbit <faerbit@posteo.net>
Date: Sun, 3 Nov 2024 23:03:03 +0100
Subject: [PATCH] Refactor stuff

---
 main.py | 55 +++++++++++++++++++++++++++++++------------------------
 1 file changed, 31 insertions(+), 24 deletions(-)

diff --git a/main.py b/main.py
index 3b1b99b..9ca02ae 100644
--- a/main.py
+++ b/main.py
@@ -11,10 +11,10 @@ class IniFile:
     def __init__(self):
         self._data = {}
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str):
         return self._data[key]
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value):
         self._data[key] = value
 
     def write(self, f):
@@ -30,16 +30,20 @@ class IniFile:
                     f.write(f"{key}={value}\n")
             f.write("\n")
 
-def abort(msg):
+def abort(msg: str):
     print(msg, file=sys.stderr)
     sys.exit(1)
 
-def enforce_list(key, dict_):
+def enforce_list(key: str, dict_: dict):
     if not isinstance(dict_[key], list):
         abort(f"{key} is not a list. Offending dict: {dict_}")
 
 
-def write_build_unit(args, yaml_dict, service_name):
+def write_file(ini_file: IniFile, path: Path):
+    with open(path, "w", encoding="utf-8") as f:
+        ini_file.write(f)
+
+def write_build_unit(args, yaml_dict: dict, service_name: str):
     out_file = args.output_dir / f"{service_name}.build"
     print(f'Generating build "{service_name}" ({out_file})')
     build = yaml_dict["services"][service_name]["build"]
@@ -58,23 +62,10 @@ def write_build_unit(args, yaml_dict, service_name):
         if "dockerfile" in build:
             unit_file["Build"]["File"] = build["dockerfile"]
 
-    with open(out_file, "w", encoding="utf-8") as f:
-        unit_file.write(f)
+    write_file(unit_file, out_file)
 
 
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("compose_file")
-    parser.add_argument("-o", "--output-dir", default="/etc/containers/systemd", type=Path)
-    parser.add_argument("-u", "--user")
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-    with open(args.compose_file, encoding="utf-8") as f:
-        yaml_dict = yaml.safe_load(f)
-
+def write_network_units(args, yaml_dict):
     for network_name in yaml_dict.get("networks", []):
         out_file = args.output_dir / f"{network_name}.network"
         print(f'Generating network "{network_name}" ({out_file})')
@@ -98,9 +89,10 @@ def main():
         unit_file["Install"] = {}
         unit_file["Install"]["WantedBy"] = "default.target"
 
-        with open(out_file, "w", encoding="utf-8") as f:
-            unit_file.write(f)
+        write_file(unit_file, out_file)
 
+
+def write_service_units(args, yaml_dict):
     for service_name in yaml_dict["services"]:
         out_file = args.output_dir / f"{service_name}.container"
         print(f'Generating container "{service_name}" ({out_file})')
@@ -170,8 +162,23 @@ def main():
         unit_file["Install"] = {}
         unit_file["Install"]["WantedBy"] = "default.target"
 
-        with open(out_file, "w", encoding="utf-8") as f:
-            unit_file.write(f)
+        write_file(unit_file, out_file)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("compose_file")
+    parser.add_argument("-o", "--output-dir", default="/etc/containers/systemd", type=Path)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    with open(args.compose_file, encoding="utf-8") as f:
+        yaml_dict = yaml.safe_load(f)
+
+    write_network_units(args, yaml_dict)
+    write_service_units(args, yaml_dict)
 
     print('\n==> Remember to run "systemctl daemon-reload"')