xref: /aosp_15_r20/external/pytorch/tools/code_coverage/package/util/utils.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1from __future__ import annotations
2
3import os
4import shutil
5import sys
6import time
7from typing import Any, NoReturn
8
9from .setting import (
10    CompilerType,
11    LOG_DIR,
12    PROFILE_DIR,
13    TestList,
14    TestPlatform,
15    TestType,
16)
17
18
19def convert_time(seconds: float) -> str:
20    seconds = int(round(seconds))
21    seconds = seconds % (24 * 3600)
22    hour = seconds // 3600
23    seconds %= 3600
24    minutes = seconds // 60
25    seconds %= 60
26
27    return "%d:%02d:%02d" % (hour, minutes, seconds)
28
29
30def print_time(message: str, start_time: float, summary_time: bool = False) -> None:
31    with open(os.path.join(LOG_DIR, "log.txt"), "a+") as log_file:
32        end_time = time.time()
33        print(message, convert_time(end_time - start_time), file=log_file)
34        if summary_time:
35            print("\n", file=log_file)
36
37
38def print_log(*args: Any) -> None:
39    with open(os.path.join(LOG_DIR, "log.txt"), "a+") as log_file:
40        print(f"[LOG] {' '.join(args)}", file=log_file)
41
42
43def print_error(*args: Any) -> None:
44    with open(os.path.join(LOG_DIR, "log.txt"), "a+") as log_file:
45        print(f"[ERROR] {' '.join(args)}", file=log_file)
46
47
48def remove_file(path: str) -> None:
49    if os.path.exists(path):
50        os.remove(path)
51
52
53def remove_folder(path: str) -> None:
54    shutil.rmtree(path)
55
56
57def create_folder(*paths: Any) -> None:
58    for path in paths:
59        os.makedirs(path, exist_ok=True)
60
61
62# clean up all the files generated by coverage tool
63def clean_up() -> None:
64    # remove profile folder
65    remove_folder(PROFILE_DIR)
66    sys.exit("Clean Up Successfully!")
67
68
69def convert_to_relative_path(whole_path: str, base_path: str) -> str:
70    # ("profile/raw", "profile") -> "raw"
71    if base_path not in whole_path:
72        raise RuntimeError(base_path + " is not in " + whole_path)
73    return whole_path[len(base_path) + 1 :]
74
75
76def replace_extension(filename: str, ext: str) -> str:
77    return filename[: filename.rfind(".")] + ext
78
79
80# a file is related if it's in one of the test_list folder
81def related_to_test_list(file_name: str, test_list: TestList) -> bool:
82    for test in test_list:
83        if test.name in file_name:
84            return True
85    return False
86
87
88def get_raw_profiles_folder() -> str:
89    return os.environ.get("RAW_PROFILES_FOLDER", os.path.join(PROFILE_DIR, "raw"))
90
91
92def detect_compiler_type(platform: TestPlatform) -> CompilerType:
93    if platform == TestPlatform.OSS:
94        from package.oss.utils import (  # type: ignore[assignment, import, misc]
95            detect_compiler_type,
96        )
97
98        cov_type = detect_compiler_type()  # type: ignore[call-arg]
99    else:
100        from caffe2.fb.code_coverage.tool.package.fbcode.utils import (  # type: ignore[import]
101            detect_compiler_type,
102        )
103
104        cov_type = detect_compiler_type()
105
106    check_compiler_type(cov_type)
107    return cov_type  # type: ignore[no-any-return]
108
109
110def get_test_name_from_whole_path(path: str) -> str:
111    # code_coverage_tool/profile/merged/haha.merged -> haha
112    start = path.rfind("/")
113    end = path.rfind(".")
114    assert start >= 0 and end >= 0
115    return path[start + 1 : end]
116
117
118def check_compiler_type(cov_type: CompilerType | None) -> None:
119    if cov_type is not None and cov_type in [CompilerType.GCC, CompilerType.CLANG]:
120        return
121    raise Exception(  # noqa: TRY002
122        f"Can't parse compiler type: {cov_type}.",
123        " Please set environment variable COMPILER_TYPE as CLANG or GCC",
124    )
125
126
127def check_platform_type(platform_type: TestPlatform) -> None:
128    if platform_type in [TestPlatform.OSS, TestPlatform.FBCODE]:
129        return
130    raise Exception(  # noqa: TRY002
131        f"Can't parse platform type: {platform_type}.",
132        " Please set environment variable COMPILER_TYPE as OSS or FBCODE",
133    )
134
135
136def check_test_type(test_type: str, target: str) -> None:
137    if test_type in [TestType.CPP.value, TestType.PY.value]:
138        return
139    raise Exception(  # noqa: TRY002
140        f"Can't parse test type: {test_type}.",
141        f" Please check the type of buck target: {target}",
142    )
143
144
145def raise_no_test_found_exception(
146    cpp_binary_folder: str, python_binary_folder: str
147) -> NoReturn:
148    raise RuntimeError(
149        f"No cpp and python tests found in folder **{cpp_binary_folder} and **{python_binary_folder}**"
150    )
151