Refactor stuff

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"')