xref: /aosp_15_r20/external/pytorch/tools/setup_helpers/cmake_utils.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1"""
2This is refactored from cmake.py to avoid circular imports issue with env.py,
3which calls get_cmake_cache_variables_from_file
4"""
5
6from __future__ import annotations
7
8import re
9from typing import IO, Optional, Union
10
11
12CMakeValue = Optional[Union[bool, str]]
13
14
15def convert_cmake_value_to_python_value(
16    cmake_value: str, cmake_type: str
17) -> CMakeValue:
18    r"""Convert a CMake value in a string form to a Python value.
19
20    Args:
21      cmake_value (string): The CMake value in a string form (e.g., "ON", "OFF", "1").
22      cmake_type (string): The CMake type of :attr:`cmake_value`.
23
24    Returns:
25      A Python value corresponding to :attr:`cmake_value` with type :attr:`cmake_type`.
26    """
27
28    cmake_type = cmake_type.upper()
29    up_val = cmake_value.upper()
30    if cmake_type == "BOOL":
31        # https://gitlab.kitware.com/cmake/community/wikis/doc/cmake/VariablesListsStrings#boolean-values-in-cmake
32        return not (
33            up_val in ("FALSE", "OFF", "N", "NO", "0", "", "NOTFOUND")
34            or up_val.endswith("-NOTFOUND")
35        )
36    elif cmake_type == "FILEPATH":
37        if up_val.endswith("-NOTFOUND"):
38            return None
39        else:
40            return cmake_value
41    else:  # Directly return the cmake_value.
42        return cmake_value
43
44
45def get_cmake_cache_variables_from_file(
46    cmake_cache_file: IO[str],
47) -> dict[str, CMakeValue]:
48    r"""Gets values in CMakeCache.txt into a dictionary.
49
50    Args:
51      cmake_cache_file: A CMakeCache.txt file object.
52    Returns:
53      dict: A ``dict`` containing the value of cached CMake variables.
54    """
55
56    results = {}
57    for i, line in enumerate(cmake_cache_file, 1):
58        line = line.strip()
59        if not line or line.startswith(("#", "//")):
60            # Blank or comment line, skip
61            continue
62
63        # Almost any character can be part of variable name and value. As a practical matter, we assume the type must be
64        # valid if it were a C variable name. It should match the following kinds of strings:
65        #
66        #   USE_CUDA:BOOL=ON
67        #   "USE_CUDA":BOOL=ON
68        #   USE_CUDA=ON
69        #   USE_CUDA:=ON
70        #   Intel(R) MKL-DNN_SOURCE_DIR:STATIC=/path/to/pytorch/third_party/ideep/mkl-dnn
71        #   "OpenMP_COMPILE_RESULT_CXX_openmp:experimental":INTERNAL=FALSE
72        matched = re.match(
73            r'("?)(.+?)\1(?::\s*([a-zA-Z_-][a-zA-Z0-9_-]*)?)?\s*=\s*(.*)', line
74        )
75        if matched is None:  # Illegal line
76            raise ValueError(f"Unexpected line {i} in {repr(cmake_cache_file)}: {line}")
77        _, variable, type_, value = matched.groups()
78        if type_ is None:
79            type_ = ""
80        if type_.upper() in ("INTERNAL", "STATIC"):
81            # CMake internal variable, do not touch
82            continue
83        results[variable] = convert_cmake_value_to_python_value(value, type_)
84
85    return results
86