flags.cpp 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244
  1. // Copyright (c) 2023 Google LLC.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. #include "flags.h"
  15. #include <algorithm>
  16. #include <cerrno>
  17. #include <cstdlib>
  18. #include <cstring>
  19. #include <iostream>
  20. #include <regex>
  21. #include <string>
  22. #include <unordered_set>
  23. #include <variant>
  24. #include <vector>
  25. namespace flags {
  26. std::vector<std::string> positional_arguments;
  27. namespace {
  28. using token_t = const char*;
  29. using token_iterator_t = token_t*;
  30. // Extracts the flag name from a potential token.
  31. // This function only looks for a '=', to split the flag name from the value for
  32. // long-form flags. Returns the name of the flag, prefixed with the hyphen(s).
  33. inline std::string get_flag_name(const std::string& flag, bool is_short_flag) {
  34. if (is_short_flag) {
  35. return flag;
  36. }
  37. size_t equal_index = flag.find('=');
  38. if (equal_index == std::string::npos) {
  39. return flag;
  40. }
  41. return flag.substr(0, equal_index);
  42. }
  43. // Parse a boolean flag. Returns `true` if the parsing succeeded, `false`
  44. // otherwise.
  45. bool parse_bool_flag(Flag<bool>& flag, bool is_short_flag,
  46. const std::string& token) {
  47. if (is_short_flag) {
  48. flag.value() = true;
  49. return true;
  50. }
  51. const std::string raw_flag(token);
  52. size_t equal_index = raw_flag.find('=');
  53. if (equal_index == std::string::npos) {
  54. flag.value() = true;
  55. return true;
  56. }
  57. const std::string value = raw_flag.substr(equal_index + 1);
  58. if (value == "true") {
  59. flag.value() = true;
  60. return true;
  61. }
  62. if (value == "false") {
  63. flag.value() = false;
  64. return true;
  65. }
  66. return false;
  67. }
  68. // Parse a uint32_t flag value.
  69. bool parse_flag_value(Flag<uint32_t>& flag, const std::string& value) {
  70. std::regex unsigned_pattern("^ *[0-9]+ *$");
  71. if (!std::regex_match(value, unsigned_pattern)) {
  72. std::cerr << "'" << value << "' is not a unsigned number." << std::endl;
  73. return false;
  74. }
  75. errno = 0;
  76. char* end_ptr = nullptr;
  77. const uint64_t number = strtoull(value.c_str(), &end_ptr, 10);
  78. if (end_ptr == nullptr || end_ptr != value.c_str() + value.size() ||
  79. errno == EINVAL) {
  80. std::cerr << "'" << value << "' is not a unsigned number." << std::endl;
  81. return false;
  82. }
  83. if (errno == ERANGE || number > static_cast<size_t>(UINT32_MAX)) {
  84. std::cerr << "'" << value << "' cannot be represented as a 32bit unsigned."
  85. << std::endl;
  86. return false;
  87. }
  88. flag.value() = static_cast<uint32_t>(number);
  89. return true;
  90. }
  91. // "Parse" a string flag value (assigns it, cannot fail).
  92. bool parse_flag_value(Flag<std::string>& flag, const std::string& value) {
  93. flag.value() = value;
  94. return true;
  95. }
  96. // Parse a potential multi-token flag. Moves the iterator to the last flag's
  97. // token if it's a multi-token flag. Returns `true` if the parsing succeeded.
  98. // The iterator is moved to the last parsed token.
  99. template <typename T>
  100. bool parse_flag(Flag<T>& flag, bool is_short_flag, const char*** iterator) {
  101. const std::string raw_flag(**iterator);
  102. std::string raw_value;
  103. const size_t equal_index = raw_flag.find('=');
  104. if (is_short_flag || equal_index == std::string::npos) {
  105. if ((*iterator)[1] == nullptr) {
  106. return false;
  107. }
  108. // This is a bi-token flag. Moving iterator to the last parsed token.
  109. raw_value = (*iterator)[1];
  110. *iterator += 1;
  111. } else {
  112. // This is a mono-token flag, no need to move the iterator.
  113. raw_value = raw_flag.substr(equal_index + 1);
  114. }
  115. return parse_flag_value(flag, raw_value);
  116. }
  117. } // namespace
  118. // This is the function to expand if you want to support a new type.
  119. bool FlagList::parse_flag_info(FlagInfo& info, token_iterator_t* iterator) {
  120. bool success = false;
  121. std::visit(
  122. [&](auto&& item) {
  123. using T = std::decay_t<decltype(item.get())>;
  124. if constexpr (std::is_same_v<T, Flag<bool>>) {
  125. success = parse_bool_flag(item.get(), info.is_short, **iterator);
  126. } else if constexpr (std::is_same_v<T, Flag<std::string>>) {
  127. success = parse_flag(item.get(), info.is_short, iterator);
  128. } else if constexpr (std::is_same_v<T, Flag<uint32_t>>) {
  129. success = parse_flag(item.get(), info.is_short, iterator);
  130. } else {
  131. static_assert(always_false_v<T>, "Unsupported flag type.");
  132. }
  133. },
  134. info.flag);
  135. return success;
  136. }
  137. bool FlagList::parse(token_t* argv) {
  138. flags::positional_arguments.clear();
  139. std::unordered_set<const FlagInfo*> parsed_flags;
  140. bool ignore_flags = false;
  141. for (const char** it = argv + 1; *it != nullptr; it++) {
  142. if (ignore_flags) {
  143. flags::positional_arguments.emplace_back(*it);
  144. continue;
  145. }
  146. // '--' alone is used to mark the end of the flags.
  147. if (std::strcmp(*it, "--") == 0) {
  148. ignore_flags = true;
  149. continue;
  150. }
  151. // '-' alone is not a flag, but often used to say 'stdin'.
  152. if (std::strcmp(*it, "-") == 0) {
  153. flags::positional_arguments.emplace_back(*it);
  154. continue;
  155. }
  156. const std::string raw_flag(*it);
  157. if (raw_flag.size() == 0) {
  158. continue;
  159. }
  160. if (raw_flag[0] != '-') {
  161. flags::positional_arguments.emplace_back(*it);
  162. continue;
  163. }
  164. // Only case left: flags (long and shorts).
  165. if (raw_flag.size() < 2) {
  166. std::cerr << "Unknown flag " << raw_flag << std::endl;
  167. return false;
  168. }
  169. const bool is_short_flag = std::strncmp(*it, "--", 2) != 0;
  170. const std::string flag_name = get_flag_name(raw_flag, is_short_flag);
  171. auto needle = std::find_if(
  172. get_flags().begin(), get_flags().end(),
  173. [&flag_name](const auto& item) { return item.name == flag_name; });
  174. if (needle == get_flags().end()) {
  175. std::cerr << "Unknown flag " << flag_name << std::endl;
  176. return false;
  177. }
  178. if (parsed_flags.count(&*needle) != 0) {
  179. std::cerr << "The flag " << flag_name << " was specified multiple times."
  180. << std::endl;
  181. return false;
  182. }
  183. parsed_flags.insert(&*needle);
  184. if (!parse_flag_info(*needle, &it)) {
  185. std::cerr << "Invalid usage for flag " << flag_name << std::endl;
  186. return false;
  187. }
  188. }
  189. // Check that we parsed all required flags.
  190. for (const auto& flag : get_flags()) {
  191. if (!flag.required) {
  192. continue;
  193. }
  194. if (parsed_flags.count(&flag) == 0) {
  195. std::cerr << "Missing required flag " << flag.name << std::endl;
  196. return false;
  197. }
  198. }
  199. return true;
  200. }
  201. // Just the public wrapper around the parse function.
  202. bool Parse(const char** argv) { return FlagList::parse(argv); }
  203. } // namespace flags