xref: /aosp_15_r20/external/pytorch/c10/util/env.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1*da0073e9SAndroid Build Coastguard Worker #pragma once
2*da0073e9SAndroid Build Coastguard Worker 
3*da0073e9SAndroid Build Coastguard Worker #include <c10/util/Exception.h>
4*da0073e9SAndroid Build Coastguard Worker #include <cstdlib>
5*da0073e9SAndroid Build Coastguard Worker #include <cstring>
6*da0073e9SAndroid Build Coastguard Worker #include <optional>
7*da0073e9SAndroid Build Coastguard Worker 
8*da0073e9SAndroid Build Coastguard Worker namespace c10::utils {
9*da0073e9SAndroid Build Coastguard Worker // Reads an environment variable and returns
10*da0073e9SAndroid Build Coastguard Worker // - std::optional<true>,              if set equal to "1"
11*da0073e9SAndroid Build Coastguard Worker // - std::optional<false>,             if set equal to "0"
12*da0073e9SAndroid Build Coastguard Worker // - nullopt,   otherwise
13*da0073e9SAndroid Build Coastguard Worker //
14*da0073e9SAndroid Build Coastguard Worker // NB:
15*da0073e9SAndroid Build Coastguard Worker // Issues a warning if the value of the environment variable is not 0 or 1.
check_env(const char * name)16*da0073e9SAndroid Build Coastguard Worker inline std::optional<bool> check_env(const char* name) {
17*da0073e9SAndroid Build Coastguard Worker #ifdef _MSC_VER
18*da0073e9SAndroid Build Coastguard Worker #pragma warning(push)
19*da0073e9SAndroid Build Coastguard Worker #pragma warning(disable : 4996)
20*da0073e9SAndroid Build Coastguard Worker #endif
21*da0073e9SAndroid Build Coastguard Worker   auto envar = std::getenv(name);
22*da0073e9SAndroid Build Coastguard Worker #ifdef _MSC_VER
23*da0073e9SAndroid Build Coastguard Worker #pragma warning(pop)
24*da0073e9SAndroid Build Coastguard Worker #endif
25*da0073e9SAndroid Build Coastguard Worker   if (envar) {
26*da0073e9SAndroid Build Coastguard Worker     if (strcmp(envar, "0") == 0) {
27*da0073e9SAndroid Build Coastguard Worker       return false;
28*da0073e9SAndroid Build Coastguard Worker     }
29*da0073e9SAndroid Build Coastguard Worker     if (strcmp(envar, "1") == 0) {
30*da0073e9SAndroid Build Coastguard Worker       return true;
31*da0073e9SAndroid Build Coastguard Worker     }
32*da0073e9SAndroid Build Coastguard Worker     TORCH_WARN(
33*da0073e9SAndroid Build Coastguard Worker         "Ignoring invalid value for boolean flag ",
34*da0073e9SAndroid Build Coastguard Worker         name,
35*da0073e9SAndroid Build Coastguard Worker         ": ",
36*da0073e9SAndroid Build Coastguard Worker         envar,
37*da0073e9SAndroid Build Coastguard Worker         "valid values are 0 or 1.");
38*da0073e9SAndroid Build Coastguard Worker   }
39*da0073e9SAndroid Build Coastguard Worker   return std::nullopt;
40*da0073e9SAndroid Build Coastguard Worker }
41*da0073e9SAndroid Build Coastguard Worker } // namespace c10::utils
42