xref: /aosp_15_r20/external/pytorch/tools/setup_helpers/env.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1from __future__ import annotations
2
3import os
4import platform
5import struct
6import sys
7from itertools import chain
8from typing import cast, Iterable
9
10
11IS_WINDOWS = platform.system() == "Windows"
12IS_DARWIN = platform.system() == "Darwin"
13IS_LINUX = platform.system() == "Linux"
14
15IS_CONDA = (
16    "conda" in sys.version
17    or "Continuum" in sys.version
18    or any(x.startswith("CONDA") for x in os.environ)
19)
20CONDA_DIR = os.path.join(os.path.dirname(sys.executable), "..")
21
22IS_64BIT = struct.calcsize("P") == 8
23
24BUILD_DIR = "build"
25
26
27def check_env_flag(name: str, default: str = "") -> bool:
28    return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]
29
30
31def check_negative_env_flag(name: str, default: str = "") -> bool:
32    return os.getenv(name, default).upper() in ["OFF", "0", "NO", "FALSE", "N"]
33
34
35def gather_paths(env_vars: Iterable[str]) -> list[str]:
36    return list(chain(*(os.getenv(v, "").split(os.pathsep) for v in env_vars)))
37
38
39def lib_paths_from_base(base_path: str) -> list[str]:
40    return [os.path.join(base_path, s) for s in ["lib/x64", "lib", "lib64"]]
41
42
43# We promised that CXXFLAGS should also be affected by CFLAGS
44if "CFLAGS" in os.environ and "CXXFLAGS" not in os.environ:
45    os.environ["CXXFLAGS"] = os.environ["CFLAGS"]
46
47
48class BuildType:
49    """Checks build type. The build type will be given in :attr:`cmake_build_type_env`. If :attr:`cmake_build_type_env`
50    is ``None``, then the build type will be inferred from ``CMakeCache.txt``. If ``CMakeCache.txt`` does not exist,
51    os.environ['CMAKE_BUILD_TYPE'] will be used.
52
53    Args:
54      cmake_build_type_env (str): The value of os.environ['CMAKE_BUILD_TYPE']. If None, the actual build type will be
55        inferred.
56
57    """
58
59    def __init__(self, cmake_build_type_env: str | None = None) -> None:
60        if cmake_build_type_env is not None:
61            self.build_type_string = cmake_build_type_env
62            return
63
64        cmake_cache_txt = os.path.join(BUILD_DIR, "CMakeCache.txt")
65        if os.path.isfile(cmake_cache_txt):
66            # Found CMakeCache.txt. Use the build type specified in it.
67            from .cmake_utils import get_cmake_cache_variables_from_file
68
69            with open(cmake_cache_txt) as f:
70                cmake_cache_vars = get_cmake_cache_variables_from_file(f)
71            # Normally it is anti-pattern to determine build type from CMAKE_BUILD_TYPE because it is not used for
72            # multi-configuration build tools, such as Visual Studio and XCode. But since we always communicate with
73            # CMake using CMAKE_BUILD_TYPE from our Python scripts, this is OK here.
74            self.build_type_string = cast(str, cmake_cache_vars["CMAKE_BUILD_TYPE"])
75        else:
76            self.build_type_string = os.environ.get("CMAKE_BUILD_TYPE", "Release")
77
78    def is_debug(self) -> bool:
79        "Checks Debug build."
80        return self.build_type_string == "Debug"
81
82    def is_rel_with_deb_info(self) -> bool:
83        "Checks RelWithDebInfo build."
84        return self.build_type_string == "RelWithDebInfo"
85
86    def is_release(self) -> bool:
87        "Checks Release build."
88        return self.build_type_string == "Release"
89
90
91# hotpatch environment variable 'CMAKE_BUILD_TYPE'. 'CMAKE_BUILD_TYPE' always prevails over DEBUG or REL_WITH_DEB_INFO.
92if "CMAKE_BUILD_TYPE" not in os.environ:
93    if check_env_flag("DEBUG"):
94        os.environ["CMAKE_BUILD_TYPE"] = "Debug"
95    elif check_env_flag("REL_WITH_DEB_INFO"):
96        os.environ["CMAKE_BUILD_TYPE"] = "RelWithDebInfo"
97    else:
98        os.environ["CMAKE_BUILD_TYPE"] = "Release"
99
100build_type = BuildType()
101