Export Salt Cloud driver functions to XLSX

export-salt-cloud-driver-functions.py

#!/usr/bin/env python3

import argparse
import ast
import os
import sys

import xlsxwriter


def parse_module_dir(module_dir, module_parser):
    modules = {}
    for module in os.listdir(module_dir):
        if not module.endswith('.py') or module == '__init__.py':
            continue
        module_path = os.path.join(module_dir, module)
        result = module_parser(module_path)
        if result is not None:
            modules[module] = result
        else:
            sys.stderr.write('Empty result for module %s\n' % module_path)
    return modules


def parse(module):
    with open(module, "r") as source:
        tree = ast.parse(source.read())
    return tree


def parse_driver(driver):
    tree = parse(driver)
    functions = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
    func_list = []
    for fn in functions:
        doc = ast.get_docstring(fn)
        doc = doc.splitlines()[0] if doc else ''
        if fn.name.startswith('_'):
            continue
        func_list.append((fn.name, doc))

    assignments = [
        node.targets[0].id for node in tree.body
        if isinstance(node, ast.Assign) and
        isinstance(node.value, ast.Call) and
        isinstance(node.value.func, ast.Name) and
        node.value.func.id == 'namespaced_function'
    ]

    for assignment in assignments:
        func_list.append((assignment, 'NAMESPACED FUNCTION'))
    func_list.sort()
    return func_list


def parse_test(test):
    name = os.path.splitext(os.path.basename(test))[0]
    cloud_name = name.split('_', 1)
    if len(cloud_name) > 1:
        return cloud_name[1]


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        '--repo-dir',
        type=str,
        help='Path to Salt repo directory',
        required=True,
    )
    parser.add_argument(
        '--out',
        type=str,
        default='salt-cloud-drivers.xlsx',
        help='Name of the resulting xlsx file',
    )
    parser.add_argument(
        '--group',
        action='store_true',
        help='Group functions by name',
    )
    options = parser.parse_args()
    if not os.path.isdir(options.repo_dir):
        parser.exit('Repo should be a directory: {}'.format(options.repo_dir))

    drivers = parse_module_dir(
        os.path.join(options.repo_dir, 'salt/cloud/clouds'),
        parse_driver
    )
    drivers = dict(sorted(drivers.items(), key=lambda i: i[0]))

    unit_tests = parse_module_dir(
        os.path.join(options.repo_dir, 'tests/unit/cloud/clouds'),
        parse_test
    )
    int_tests = parse_module_dir(
        os.path.join(options.repo_dir, 'tests/integration/cloud/clouds'),
        parse_test
    )
    workbook = xlsxwriter.Workbook(options.out)
    worksheet = workbook.add_worksheet()

    def make_header(drv_name):
        header = os.path.splitext(drv_name)[0]
        if header in unit_tests.values():
            header += ' +unit'
        if header in int_tests.values():
            header += ' +int'
        return header

    if options.group:
        all_funcs = sorted(list(set([f[0] for k, v in drivers.items() for f in v])))

        for fi, fn in enumerate(all_funcs):
            worksheet.write(fi + 1, 0, fn)

        for drv_index, (drv_name, functions) in enumerate(drivers.items()):
            worksheet.write(0, drv_index + 1, make_header(drv_name))
            for fi, fn in enumerate(functions):
                fn_index = all_funcs.index(fn[0])
                worksheet.write(fn_index + 1, drv_index + 1, fn[0])
                worksheet.write_comment(fn_index + 1, drv_index + 1, fn[1])
    else:
        for drv_index, (drv_name, functions) in enumerate(drivers.items()):
            worksheet.write(0, drv_index + 1, make_header(drv_name))
            for fi, fn in enumerate(functions):
                worksheet.write(fi + 1, drv_index + 1, fn[0])
                worksheet.write_comment(fi + 1, drv_index + 1, fn[1])

    workbook.close()