monakhos/ansible/parse_arch_packages.py

import argparse
import json
import sys
import yaml

def flatten_packages(item):
    packages = []
    if isinstance(item, str):
        return {'name': item}
    if isinstance(item, dict) and 'name' in item:
        return item
    if isinstance(item, list):
        return [flatten_packages(p) for p in item]
    if isinstance(item, dict):
        packages = []
        for v in item.values():
            packages += flatten_packages(v)
        return packages
    return []

def output_packages(packages_tree, aur_filter: bool = False):
    flat_packages = flatten_packages(packages_tree)

    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

def cli():
    parser = argparse.ArgumentParser(description='Optional app description')
    parser.add_argument('categories', type=str, help="Categories to include, all for all")
    # for now, category must refer to first level nodes
    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()