microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
base/env_string.h
105lines · modeblame
22340011cao lei2 years ago | 1 | // Copyright (c) Microsoft Corporation. All rights reserved. |
| 2 | // Licensed under the MIT License. | |
| 3 | #pragma once | |
| 4 | #include <locale> | |
| 5 | #include <optional> | |
| 6 | #include <string> | |
| 7 | #include <sstream> | |
97ee9eb5Wenbing Li2 years ago | 8 | |
22340011cao lei2 years ago | 9 | #include "string_utils.h" |
| 10 | #ifdef _WIN32 | |
| 11 | #include <Windows.h> | |
| 12 | #endif | |
| 13 | | |
| 14 | template <typename T> | |
| 15 | bool TryParseStringWithClassicLocale(std::string_view str, T& value) { | |
| 16 | if constexpr (std::is_integral<T>::value && std::is_unsigned<T>::value) { | |
| 17 | // if T is unsigned integral type, reject negative values which will wrap | |
| 18 | if (!str.empty() && str[0] == '-') { | |
| 19 | return false; | |
| 20 | } | |
| 21 | } | |
| 22 | | |
| 23 | // don't allow leading whitespace | |
| 24 | if (!str.empty() && std::isspace(str[0], std::locale::classic())) { | |
| 25 | return false; | |
| 26 | } | |
| 27 | | |
| 28 | std::istringstream is{std::string{str}}; | |
| 29 | is.imbue(std::locale::classic()); | |
| 30 | T parsed_value{}; | |
| 31 | | |
| 32 | const bool parse_successful = | |
| 33 | is >> parsed_value && | |
| 34 | is.get() == std::istringstream::traits_type::eof(); // don't allow trailing characters | |
| 35 | if (!parse_successful) { | |
| 36 | return false; | |
| 37 | } | |
| 38 | | |
| 39 | value = std::move(parsed_value); | |
| 40 | return true; | |
| 41 | } | |
| 42 | | |
| 43 | inline bool TryParseStringWithClassicLocale(std::string_view str, std::string& value) { | |
| 44 | value = str; | |
| 45 | return true; | |
| 46 | } | |
| 47 | | |
| 48 | inline bool TryParseStringWithClassicLocale(std::string_view str, bool& value) { | |
| 49 | if (str == "0" || str == "False" || str == "false") { | |
| 50 | value = false; | |
| 51 | return true; | |
| 52 | } | |
| 53 | | |
| 54 | if (str == "1" || str == "True" || str == "true") { | |
| 55 | value = true; | |
| 56 | return true; | |
| 57 | } | |
| 58 | | |
| 59 | return false; | |
| 60 | } | |
| 61 | | |
| 62 | template <typename T> | |
| 63 | std::optional<T> ParseEnvironmentVariable(const std::string& name) { | |
| 64 | std::string buffer; | |
| 65 | #ifdef _WIN32 | |
| 66 | constexpr size_t kBufferSize = 32767; | |
| 67 | | |
| 68 | // Create buffer to hold the result | |
| 69 | buffer.resize(kBufferSize, '\0'); | |
| 70 | | |
| 71 | // The last argument is the size of the buffer pointed to by the lpBuffer parameter, including the null-terminating character, in characters. | |
| 72 | // If the function succeeds, the return value is the number of characters stored in the buffer pointed to by lpBuffer, not including the terminating null character. | |
| 73 | // Therefore, If the function succeeds, kBufferSize should be larger than char_count. | |
| 74 | auto char_count = GetEnvironmentVariableA(name.c_str(), buffer.data(), kBufferSize); | |
| 75 | | |
| 76 | if (kBufferSize > char_count) { | |
| 77 | buffer.resize(char_count); | |
| 78 | } else { | |
| 79 | // Else either the call was failed, or the buffer wasn't large enough. | |
| 80 | // TODO: Understand the reason for failure by calling GetLastError(). | |
| 81 | // If it is due to the specified environment variable being found in the environment block, | |
| 82 | // GetLastError() returns ERROR_ENVVAR_NOT_FOUND. | |
| 83 | // For now, we assume that the environment variable is not found. | |
| 84 | buffer.clear(); | |
| 85 | } | |
| 86 | #else | |
| 87 | char* val = getenv(name.c_str()); | |
| 88 | buffer = (val == nullptr) ? std::string() : std::string(val); | |
| 89 | #endif | |
| 90 | T parsed_value; | |
| 91 | if (!TryParseStringWithClassicLocale(buffer, parsed_value)) { | |
| 92 | OrtW::Exception(MakeString("Failed to parse environment variable - name: ", name, ", value: ", buffer), OrtErrorCode::ORT_FAIL); | |
| 93 | } | |
| 94 | return parsed_value; | |
| 95 | } | |
| 96 | | |
| 97 | template <typename T> | |
| 98 | T ParseEnvironmentVariableWithDefault(const std::string& name, const T& default_value) { | |
| 99 | const auto parsed = ParseEnvironmentVariable<T>(name); | |
| 100 | if (parsed.has_value()) { | |
| 101 | return *parsed; | |
| 102 | } | |
| 103 | | |
| 104 | return default_value; | |
| 105 | } |