refactor(packages): directly pipe python script into ansible

ansible/parse_arch_packages.py

@@ -1,13 +1,7 @@
+import argparse
 import json
 import sys
 import yaml
-import fileinput
-
-packages_tree = None
-yaml_config = ''.join(sys.stdin.readlines())
-packages_tree = yaml.safe_load(yaml_config)
-
-assert packages_tree is not None, "Must have load valid package tree"
 
 def flatten_packages(item):
     packages = []
@@ -19,24 +13,46 @@ def flatten_packages(item):
         return [flatten_packages(p) for p in item]
     if isinstance(item, dict):
         packages = []
-        for k, v in item.items():
+        for v in item.values():
             packages += flatten_packages(v)
         return packages
     return []
 
-flat_packages = flatten_packages(packages_tree)
+def output_packages(packages_tree, aur_filter: bool = False):
+    flat_packages = flatten_packages(packages_tree)
 
-native_names = []
-aur_names = []
-for p in flat_packages:
-    assert isinstance(p, dict)
-    if p['name'].startswith('aur/'):
-        aur_names.append(p['name'].split('/')[1])
-        continue
-    native_names.append(p['name'])
+    names = []
+    for p in flat_packages:
+        assert isinstance(p, dict)
+        is_aur = p['name'].startswith('aur/')
+        if (aur_filter and not is_aur) or (not aur_filter and is_aur):
+            continue
+        names.append(
+            p['name'].split('/')[1] if '/' in p['name'] else p['name']
+        )
+    return names
 
-print(json.dumps({
-    "native": native_names,
-    "aur": aur_names
-}))
-# print(' '.join(names))
+def cli():
+    parser = argparse.ArgumentParser(description='Optional app description')
+    parser.add_argument('categories', type=str, help="Categories to include, all for all")
+    parser.add_argument('--aur', action='store_true', help="Also output AUR packages")
+    args = parser.parse_args()
+
+    # load data from stdin
+    packages_tree = None
+    yaml_config = ''.join(sys.stdin.readlines())
+    packages_tree = yaml.safe_load(yaml_config)
+    assert packages_tree is not None, "Must have load valid package tree"
+
+    tree: dict = {}
+    if args.categories == 'all':
+        tree = packages_tree
+    else:
+        selected_categories = args.categories.split(',')
+        for category in selected_categories:
+            if category not in packages_tree:
+                continue
+            tree = tree | {category: packages_tree[category]}
+    print(json.dumps(output_packages(tree, args.aur)))
+
+cli()