Add choices to getopt options

3 files changed, 60 insertions, 11 deletions

diff --git a/library/cpp/getopt/last_getopt_demo/demo.cpp b/library/cpp/getopt/last_getopt_demo/demo.cpp
index 79426a9cc9..a0e82a936c 100644
--- a/library/cpp/getopt/last_getopt_demo/demo.cpp
+++ b/library/cpp/getopt/last_getopt_demo/demo.cpp
@@ -81,17 +81,16 @@ protected:
             .Help("specify HTTP method")
             .CompletionArgHelp("http method")
             .StoreResult(&ExplicitMethod_)
-            .Completer(
-                NLastGetopt::NComp::Choice(
-                    {{"GET", "request representation of the specified resource"},
-                     {"HEAD", "request response identical to that of GET, but without response body"},
-                     {"POST", "submit an entry to the specified resource"},
-                     {"PUT", "replace representation of the specified resource with the request body"},
-                     {"DELETE", "delete the specified resource"},
-                     {"CONNECT", "establish a tunnel to the server identified by the target resource"},
-                     {"OPTIONS", "describe the communication options for the target resource"},
-                     {"TRACE", "perform a message loop-back test"},
-                     {"PATCH", "apply partial modifications to the specified resource"}}));
+            .ChoicesWithCompletion({
+                {"GET", "request representation of the specified resource"},
+                {"HEAD", "request response identical to that of GET, but without response body"},
+                {"POST", "submit an entry to the specified resource"},
+                {"PUT", "replace representation of the specified resource with the request body"},
+                {"DELETE", "delete the specified resource"},
+                {"CONNECT", "establish a tunnel to the server identified by the target resource"},
+                {"OPTIONS", "describe the communication options for the target resource"},
+                {"TRACE", "perform a message loop-back test"},
+                {"PATCH", "apply partial modifications to the specified resource"}});
 
         opts.AddLongOption('U', "user-agent")
             .RequiredArgument("agent-string")
diff --git a/library/cpp/getopt/small/last_getopt_opt.h b/library/cpp/getopt/small/last_getopt_opt.h
index 4ffdfc394a..32afcb0437 100644
--- a/library/cpp/getopt/small/last_getopt_opt.h
+++ b/library/cpp/getopt/small/last_getopt_opt.h
@@ -4,11 +4,13 @@
 #include "last_getopt_handlers.h"
 
 #include <util/string/split.h>
+#include <util/generic/hash_set.h>
 #include <util/generic/ptr.h>
 #include <util/generic/string.h>
 #include <util/generic/maybe.h>
 #include <util/generic/vector.h>
 #include <util/string/cast.h>
+#include <util/string/join.h>
 
 #include <optional>
 #include <stdarg.h>
@@ -80,6 +82,7 @@ namespace NLastGetopt {
         TdOptVal OptionalValue_;
         TdOptVal DefaultValue_;
         TOptHandlers Handlers_;
+        THashSet<TString> Choices_;
 
     public:
         /**
@@ -398,6 +401,10 @@ namespace NLastGetopt {
             return Help_;
         }
 
+        TString GetChoicesHelp() const {
+            return JoinSeq(", ", Choices_);
+        }
+
         /**
          * Set help string that appears when argument completer lists available options.
          *
@@ -728,6 +735,41 @@ namespace NLastGetopt {
         TOpt& KVHandler(TpFunc func, const char kvdelim = '=') {
             return Handler(new NLastGetopt::TOptKVHandler<TpFunc>(func, kvdelim));
         }
+
+        template <typename TIterator>
+        TOpt& Choices(TIterator begin, TIterator end) {
+            return Choices(THashSet<typename TIterator::value_type>{begin, end});
+        }
+
+        template <typename TValue>
+        TOpt& Choices(THashSet<TValue> choices) {
+            Choices_ = std::move(choices);
+            return Handler1T<TValue>(
+                [this] (const TValue& arg) {
+                    if (!Choices_.contains(arg)) {
+                        throw TUsageException() << " value '" << arg
+                                                << "' is not allowed for option '" << GetName() << "'";
+                    }
+                });
+        }
+
+        TOpt& Choices(TVector<TString> choices) {
+            return Choices(
+                THashSet<TString>{
+                    std::make_move_iterator(choices.begin()),
+                    std::make_move_iterator(choices.end())
+                });
+        }
+
+        TOpt& ChoicesWithCompletion(TVector<NComp::TChoice> choices) {
+            Completer(NComp::Choice(choices));
+            THashSet<TString> choicesSet;
+            choicesSet.reserve(choices.size());
+            for (const auto& choice : choices) {
+                choicesSet.insert(choice.Choice);
+            }
+            return Choices(std::move(choicesSet));
+        }
     };
 
     /**
diff --git a/library/cpp/getopt/small/last_getopt_opts.cpp b/library/cpp/getopt/small/last_getopt_opts.cpp
index 03c432849f..939d5d59d2 100644
--- a/library/cpp/getopt/small/last_getopt_opts.cpp
+++ b/library/cpp/getopt/small/last_getopt_opts.cpp
@@ -427,6 +427,14 @@ namespace NLastGetopt {
                     os << Wrap(Wrap_, help, SPad + leftPadding + " ", &lastLineLength, &helpHasParagraphs);
                 }
 
+                auto choicesHelp = opt->GetChoicesHelp();
+                if (!choicesHelp.empty()) {
+                    if (help) {
+                        os << Endl << SPad << leftPadding << " ";
+                    }
+                    os << "(values: " << choicesHelp << ")";
+                }
+
                 if (opt->HasDefaultValue()) {
                     auto quotedDef = QuoteForHelp(opt->GetDefaultValue());
                     if (helpHasParagraphs) {