llvm-project/flang/examples/FlangOmpReport/yaml_summarizer.py

"""YAML Summariser

The flang plugin ``flang-omp-report`` takes one fortran
file in and returns a YAML report file of the input file.
This becomes an issue when you want to analyse an entire project
into one final report.
The purpose of this Python script is to generate a final YAML
summary from all of the files generated by ``flang-omp-report``.

Currently, it requires ``ruamel.yaml``,
which can be installed with:

    ``pip3 install ruamel.yaml``

By default it scans the directory it is ran in
for any YAML files and outputs a summary to
stdout. It can be ran as:

    ``python3 yaml_summarizer.py``

Parameters:

    -d   --directory   Specify which directory to scan. Multiple directories can be searched by
                       providing a semicolon seperated list of directories.

    -l   --log         Combine all yaml files into one log (instead of generating a summary)

    -o   --output      Specify a directory in which to save the summary file

    -r   --recursive   Recursively search directory for all yaml files

Examples:

    ``python3 yaml_summarizer.py -d ~/llvm-project/build/ -r``

    ``python3 yaml_summarizer.py -d "~/llvm-project/build/;~/llvm-project/flang/test/Examples"``

    ``python3 yaml_summarizer.py -l -o ~/examples/report.yaml``

Pseudo-examples:

    Summary:

    $ python3 yaml_summarizer.py file_1.yaml file_2.yaml
    <Unique OMP constructs with there grouped clauses from file_1.yaml and file_2.yaml>

    Construcsts are in the form:
    - construct: someOMPconstruct
    count: 8
    clauses:
    - clause: clauseOne
        count: 4
    - clause: ClauseTwo
        count: 2

    Log:

    $ python3 yaml_summarizer.py -l file_1.yaml file_2.yaml
    file_1.yaml
    <OMP clauses and constructs from file_1.yaml>
    file_2.yaml
    <OMP clauses and constructs from file_2.yaml>

    Constructs are in the form:
    - construct: someOMPConstruct
    line: 12
    clauses:
    - clause: clauseOne
        details: 'someDetailForClause'
"""

import sys
import glob
import argparse
from pathlib import Path
from os.path import isdir

from ruamel.yaml import YAML

def find_yaml_files(search_directory: Path, search_pattern: str):
    """
    Find all '.yaml' files and returns an iglob iterator to them.

    Keyword arguments:
    search_pattern -- Search pattern for 'iglob' to use for finding '.yaml' files.
                      If this is set to 'None', then it will default to just searching
                      for all '.yaml' files in the current directory.
    """
    # @TODO: Currently *all* yaml files are read - regardless of whether they have
    # been generated with  'flang-omp-report' or not. This might result in the script
    # reading files that it should ignore.
    if search_directory:
        return glob.iglob(str(search_directory.joinpath(search_pattern)), recursive=True)

    return glob.iglob(str("/" + search_pattern), recursive=True)

def process_log(data, result: list):
    """
    Process the data input as a 'log' to the result array. This esssentially just
    stitches together all of the input '.yaml' files into one result.

    Keyword arguments:
    data -- Data from yaml.load() for a yaml file. So the type can be 'Any'.
    result -- Array to add the processed data to.
    """
    for datum in data:
        items = result.get(datum['file'], [])
        items.append({"construct" : datum['construct'],
                        "line" : datum['line'],
                        "clauses" : datum['clauses']})
        result[datum['file']] = items

def add_clause(datum, construct):
    """
    Add clauses to the construct if they're missing
    Otherwise increment their count by one.

    Keyword arguments:
    datum -- Data construct containing clauses to check.
    construct -- Construct to add or increment clause count.
    """
    to_check = [i['clause'] for i in construct['clauses']]
    to_add = [i['clause'] for i in datum['clauses']]
    clauses = construct["clauses"]
    for item in to_add:
        if item in to_check:
            for clause in clauses:
                if clause["clause"] == item:
                    clause["count"] += 1
        else:
            clauses.append({"clause" : item,
                            "count" : 1})

def process_summary(data, result: dict):
    """
    Process the data input as a 'summary' to the 'result' dictionary.

    Keyword arguments:
    data -- Data from yaml.load() for a yaml file. So the type can be 'Any'.
    result -- Dictionary to add the processed data to.
    """
    for datum in data:
        construct = next((item for item in result
                            if item["construct"] == datum["construct"]), None)
        clauses = []
        # Add the construct and clauses to the summary if
        # they haven't been seen before
        if not construct:
            for i in datum['clauses']:
                clauses.append({"clause" : i['clause'],
                                "count"    : 1})
            result.append({"construct" : datum['construct'],
                            "count" : 1,
                            "clauses" : clauses})
        else:
            construct["count"] += 1

            add_clause(datum, construct)

def clean_summary(result):
    """ Cleans the result after processing the yaml files with summary format."""
    # Remove all "clauses" that are empty to keep things compact
    for construct in result:
        if construct["clauses"] == []:
            construct.pop("clauses")

def clean_log(result):
    """ Cleans the result after processing the yaml files with log format."""
    for constructs in result.values():
        for construct in constructs:
            if construct["clauses"] == []:
                construct.pop("clauses")

def output_result(yaml: YAML, result, output_file: Path):
    """
    Outputs result to either 'stdout' or to a output file.

    Keyword arguments:
    result -- Format result to output.
    output_file -- File to output result to. If this is 'None' then result will be
                   outputted to 'stdout'.
    """
    if output_file:
        with open(output_file, 'w+', encoding='utf-8') as file:
            if output_file.suffix == ".yaml":
                yaml.dump(result, file)
            else:
                file.write(result)
    else:
        yaml.dump(result, sys.stdout)

def process_yaml(search_directories: list, search_pattern: str,
                 result_format: str, output_file: Path):
    """
    Reads each yaml file, calls the appropiate format function for
    the file and then ouputs the result to either 'stdout' or to an output file.

    Keyword arguments:
    search_directories -- List of directory paths to search for '.yaml' files in.
    search_pattern -- String pattern formatted for use with glob.iglob to find all
                      '.yaml' files.
    result_format -- String representing output format. Current supported strings are: 'log'.
    output_file -- Path to output file (If value is None, then default to outputting to 'stdout').
    """
    if result_format == "log":
        result = {}
        action = process_log
        clean_report = clean_log
    else:
        result = []
        action = process_summary
        clean_report = clean_summary

    yaml = YAML()

    for search_directory in search_directories:
        for file in find_yaml_files(search_directory, search_pattern):
            with open(file, "r", encoding='utf-8') as yaml_file:
                data = yaml.load(yaml_file)
                action(data, result)

    if clean_report is not None:
        clean_report(result)

    output_result(yaml, result, output_file)

def create_arg_parser():
    """ Create and return a argparse.ArgumentParser modified for script. """
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--directory", help="Specify a directory to scan",
                        dest="dir", type=str)
    parser.add_argument("-o", "--output", help="Writes to a file instead of\
                                                stdout", dest="output", type=str)
    parser.add_argument("-r", "--recursive", help="Recursive search for .yaml files",
                        dest="recursive", type=bool, nargs='?', const=True, default=False)

    exclusive_parser = parser.add_mutually_exclusive_group()
    exclusive_parser.add_argument("-l", "--log", help="Modifies report format: "
                                  "Combines the log '.yaml' files into one file.",
                                  action='store_true', dest='log')
    return parser

def parse_arguments():
    """ Parses arguments given to script and returns a tuple of processed arguments. """
    parser = create_arg_parser()
    args = parser.parse_args()

    if args.dir:
        search_directory = [Path(path) for path in args.dir.split(";")]
    else:
        search_directory = [Path.cwd()]

    if args.recursive:
        search_pattern = "**/*.yaml"
    else:
        search_pattern = "*.yaml"

    if args.log:
        result_format = "log"
    else:
        result_format = "summary"

    if args.output:
        if isdir(args.output):
            output_file = Path(args.output).joinpath("summary.yaml")
        elif isdir(Path(args.output).resolve().parent):
            output_file = Path(args.output)
    else:
        output_file = None

    return (search_directory, search_pattern, result_format, output_file)

def main():
    """ Main function of script. """
    (search_directory, search_pattern, result_format, output_file) = parse_arguments()

    process_yaml(search_directory, search_pattern, result_format, output_file)

    return 0

if __name__ == "__main__":
    sys.exit(main())