create_benchmark.py

#!/usr/bin/env python3

import argparse
import os
import re


def main():
    p = argparse.ArgumentParser()
    p.add_argument("name", help="The name of the new benchmark to be created")
    args = p.parse_args()

    # adds benchmark to `CMakeLists.txt`
    update_cmakelists(args.name)
    # create benchmark Swift file
    create_benchmark_file(args.name)
    # imports the benchmark module in `main.swift`
    add_import_benchmark(args.name)
    # registers the benchmark with the driver in `main.swift`
    add_register_benchmark(args.name)


def update_cmakelists(name):
    """Adds a new entry to the `CMakeLists.txt` file with the given
    benchmark name.
    """
    relative_path = create_relative_path("../CMakeLists.txt")

    file_contents = []
    with open(relative_path, "r") as f:
        file_contents = f.readlines()

    file_new_contents = insert_line_alphabetically(
        name,
        "    single-source/" + name + "\n",
        file_contents,
        r"    single-source\/([a-zA-Z]+)",
    )
    with open(relative_path, "w") as f:
        for line in file_new_contents:
            f.write(line)


def create_benchmark_file(name):
    """Creates a new Swift file with the given name based on the template
    and places it in the `single-source` directory.
    """

    template_path = create_relative_path("Template.swift")
    benchmark_template = ""
    with open(template_path, "r") as f:
        benchmark_template = "".join(f.readlines())

    # fill in template with benchmark name.
    formatted_template = benchmark_template.format(name=name)

    relative_path = create_relative_path("../single-source/")
    source_file_path = os.path.join(relative_path, name + ".swift")
    with open(source_file_path, "w") as f:
        f.write(formatted_template)


def add_import_benchmark(name):
    """Adds an `import` statement to the `main.swift` file for the new
    benchmark.
    """
    relative_path = create_relative_path("../utils/main.swift")

    # read current contents into an array
    file_contents = []
    with open(relative_path, "r") as f:
        file_contents = f.readlines()

    # the test dependencies are placed before all benchmarks, so we have to
    # insert the benchmark in the right alphabetical order after we have seen
    # all test dependencies.
    read_test_dependencies = False
    previous_benchmark_name = None
    file_new_contents = []
    for line in file_contents:
        # check if this line is a definition of a benchmark and get its name
        match = re.search(r"import ([a-zA-Z]+)", line)
        if match and match.group(1):
            benchmark_name = match.group(1)
            # find where to insert the new benchmark in the right alphabetical
            # order.
            if (
                name < benchmark_name
                and previous_benchmark_name is None
                or name < benchmark_name
                and name > previous_benchmark_name
            ):
                if read_test_dependencies:
                    file_new_contents.append("import " + name + "\n" + line)
                else:
                    # all test dependencies are first specified, so from now
                    # on we can look where to insert the new benchmark.
                    read_test_dependencies = True
                    file_new_contents.append(line)
            else:
                file_new_contents.append(line)
            previous_benchmark_name = benchmark_name
        else:
            file_new_contents.append(line)
    with open(relative_path, "w") as f:
        for line in file_new_contents:
            f.write(line)


def add_register_benchmark(name):
    """Adds an `import` statement to the `main.swift` file for the new
    benchmark.
    """
    relative_path = create_relative_path("../utils/main.swift")

    file_contents = []
    with open(relative_path, "r") as f:
        file_contents = f.readlines()

    file_new_contents = insert_line_alphabetically(
        name,
        "registerBenchmark(" + name + ")\n",
        file_contents,
        r"registerBenchmark\(([a-zA-Z]+)\)",
    )
    with open(relative_path, "w") as f:
        for line in file_new_contents:
            f.write(line)


def insert_line_alphabetically(name, new_line, lines, regex):
    """Iterates through the given lines and executes the regex on each line to
    find where the new benchmark should be inserted with the given `new_line`.
    """
    # the name of the previous seen benchmark in order to insert the new
    # one at the correct position
    previous_benchmark_name = None
    # the new contents of the file
    updated_lines = []
    for line in lines:
        # apply regex and get name of benchmark on this line
        match = re.search(regex, line)
        if match and match.group(1):
            benchmark_name = match.group(1)
            # check if we're at the line where we have to insert the new
            # benchmark in the correct alphabetical order
            if (
                name < benchmark_name
                and previous_benchmark_name is None
                or name < benchmark_name
                and name > previous_benchmark_name
            ):
                updated_lines.append(new_line + line)
            else:
                updated_lines.append(line)
            previous_benchmark_name = benchmark_name
        else:
            updated_lines.append(line)
    return updated_lines


def create_relative_path(file_path):
    return os.path.join(os.path.dirname(__file__), file_path)


if __name__ == "__main__":
    main()