pytorch

1
"""
2
This is refactored from cmake.py to avoid circular imports issue with env.py,
3
which calls get_cmake_cache_variables_from_file
4
"""
5

6
from __future__ import annotations
7

8
import re
9
from typing import IO, Optional, Union
10

11

12
CMakeValue = Optional[Union[bool, str]]
13

14

15
def 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

45
def 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