Add choices to getopt options

2 files changed, 50 insertions, 0 deletions

diff --git a/library/cpp/getopt/small/last_getopt_opt.h b/library/cpp/getopt/small/last_getopt_opt.h
index 4ffdfc394a1..32afcb04377 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 03c432849f1..939d5d59d26 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) {