pytorch/tools/code_coverage/package/util/utils_init.py

import argparse
import os
from typing import Any

from .setting import (
    JSON_FOLDER_BASE_DIR,
    LOG_DIR,
    MERGED_FOLDER_BASE_DIR,
    PROFILE_DIR,
    SUMMARY_FOLDER_DIR,
    Option,
)
from .utils import create_folder, get_raw_profiles_folder, remove_file


def remove_files() -> None:
    # remove log
    remove_file(os.path.join(LOG_DIR, "log.txt"))


def create_folders() -> None:
    create_folder(
        PROFILE_DIR,
        MERGED_FOLDER_BASE_DIR,
        JSON_FOLDER_BASE_DIR,
        get_raw_profiles_folder(),
        SUMMARY_FOLDER_DIR,
        LOG_DIR,
    )


def add_arguments_utils(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
    parser.add_argument("--run", help="run the cpp test binaries", action="store_true")
    parser.add_argument(
        "--merge",
        help="merge raw profiles (only apply to clang coverage)",
        action="store_true",
    )
    parser.add_argument(
        "--export", help="generate json report for each file", action="store_true"
    )
    parser.add_argument(
        "--summary",
        help="read json report and generate file/line-oriented summary",
        action="store_true",
    )
    parser.add_argument(
        "--interested-folder",
        help="Final report will be only about these folders and its sub-folders; for example: caff2/c10;",
        nargs="+",
        default=None,
    )
    parser.add_argument(
        "--clean",
        help="delete all files generated by coverage tool",
        action="store_true",
        default=False,
    )

    return parser


def have_option(have_stage: bool, option: int) -> int:
    if have_stage:
        return option
    else:
        return 0


def get_options(args: Any) -> Option:
    option: Option = Option()
    if args.__contains__("build"):
        if args.build:
            option.need_build = True

    if args.__contains__("run"):
        if args.run:
            option.need_run = True

    if args.__contains__("merge"):
        if args.merge:
            option.need_merge = True

    if args.__contains__("export"):
        if args.export:
            option.need_export = True

    if args.__contains__("summary"):
        if args.summary:
            option.need_summary = True

    # user does not have specified stage like run
    if not any(vars(option).values()):
        option.need_build = True
        option.need_run = True
        option.need_merge = True
        option.need_export = True
        option.need_summary = True
        option.need_pytest = True

    return option