microsoft/onnxruntime-extensions
Publicmirrored fromhttps://github.com/microsoft/onnxruntime-extensionsAvailable
base/env_string.h
105lines · modecode
| 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> |
| 8 | |
| 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 | } |
| 106 | |