public-types-tests.py

#!/usr/bin/env python3

from subprocess import run
import argparse
import datetime
import json
import os

PREAMBLE = """
// DO NOT EDIT. Automatically generated by 'scripts/public-types-tests.py'
// on {date}.
""".lstrip()

CMD = [
    "cargo",
    "rustdoc",
    "-q",
    "--",
    "-Z",
    "unstable-options",
    "--output-format",
    "json",
]

def gen_docs():
    """Runs the command to retrieve the docs of this crate.

    If the command fails, the script exits.
    """
    doc_gen_result = run(CMD)
    if doc_gen_result.returncode != 0:
        print("Unable to generate the docs for the regex")
        exit(1)


def get_public_of_kind(kinds, elements, docs):
    """Gets the items that are of some kind in the public API.

    `kinds` is a list of item kinds. Ex: ["module", "struct"]
    `elements` are the keys in the docs json schema to search in.
    `docs` is usually the slice of the json tree that we care about.
    """
    keys_of_kind = filter(lambda key: docs[key]["kind"] in kinds, elements)
    public_of_kind = filter(
        lambda key: docs[key]["visibility"] == "public", keys_of_kind
    )
    return map(lambda key: docs[key], public_of_kind)


def gen_test(path):
    """Generates 4 valid lines of Rust code that assert marker traits on a type.

    `path` is a fully qualified type. Ex: regex::bytes::Regex
    """
    tests = [
        f"\tassert_send::<{path}>();",
        f"\tassert_sync::<{path}>();",
        f"\tassert_unwind_safe::<{path}>();",
        f"\tassert_ref_unwind_safe::<{path}>();",
        "",
    ]
    return "\n".join(tests)


def get_public_types():
    """Reads the public types from the json dump of the docs of this crate.

    After generating the docs, this function is used to read the dump in
    json format located at target/doc/regex.json. Then, it generates the
    fully qualified paths for each type of the public API of the crate.
    """
    with open(os.path.join("target", "doc", "regex.json")) as docs_file:
        docs = json.load(docs_file)
        index = docs["index"]

        type_paths = []
        public_modules = get_public_of_kind(["module"], index.keys(), index)
        for mod in public_modules:
            types = get_public_of_kind(["struct"], mod["inner"]["items"], index)
            for type in types:
                prefix = "" if mod["inner"]["is_crate"] else "regex::"
                type_path = f"{prefix}{mod['name']}::{type['name']}"
                type_paths.append(type_path)

        return type_paths


def gen_tests():
    tests = []
    types = get_public_types()
    for type in types:
        tests.append("")
        test = gen_test(type)
        tests.append(test)

    return tests


def get_assert_definitions():
    """Generates the definitions of the assert functions we use to test marker
    traits."""
    definitions = [
        "\tfn assert_send<T: Send>() {}",
        "\tfn assert_sync<T: Sync>() {}",
        "\tfn assert_unwind_safe<T: UnwindSafe>() {}",
        "\tfn assert_ref_unwind_safe<T: RefUnwindSafe>() {}",
    ]
    return "\n".join(definitions)


def write_tests():
    """Writes the generated tests to the tests/marker_traits.rs file."""
    with open(os.path.join("tests", "marker_traits.rs"), "w") as f:
        f.write(PREAMBLE.format(date=str(datetime.datetime.now())))
        f.write("\n")
        f.write("#[test]\nfn marker_traits() {\n")
        f.write("\tuse std::panic::{RefUnwindSafe, UnwindSafe};\n")
        f.write("\n")
        f.write(get_assert_definitions())
        f.write("\n")
        for test in gen_tests():
            f.write(f"{test}\n")
        f.write("}\n")


def check_existence():
    """Fails the script when a type in the API doesn't have a corresponding marker trait test
    in the tests/marker_traits.rs file."""
    types = get_public_types()
    test_path = os.path.join("tests", "marker_traits.rs")
    with open(test_path ) as f:
        lines = f.readlines()
        for type in types:
            if any(type in line for line in lines):
                continue
            else:
                print("Failed check!\n")
                print(f"A marker trait test for {type} was not found in {test_path}.")
                exit(1)

    print("Success!\nThe existence check found that every type has a corresponding test.")
    exit(0)


def main():
    p = argparse.ArgumentParser(
        "A script to manage marker-traits tests.",
        description="Checks for the existence of tests for the public API types, "
        "failing when there are missing tests for a marker trait of "
        "a type.",
    )
    p.add_argument(
        "-g",
        "--generate-tests",
        action="store_true",
        help=(
            "If set, the script will generate a "
            "tests/marker_traits.rs file containing a test that "
            "checks the auto traits of the public API types "
            "(These tests might need to be updated after generating them)."
        ),
    )

    args = p.parse_args()

    gen_docs()

    if args.generate_tests:
        write_tests()
    else:
        check_existence()


if __name__ == "__main__":
    main()