intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 367 insertions, 0 deletions

diff --git a/library/cpp/getopt/small/completer.cpp b/library/cpp/getopt/small/completer.cpp
new file mode 100644
index 0000000000..3fff684adb
--- /dev/null
+++ b/library/cpp/getopt/small/completer.cpp
@@ -0,0 +1,367 @@
+#include "completer.h"
+
+#include "completion_generator.h"
+
+#include <util/string/cast.h>
+#include <util/generic/fwd.h>
+
+using NLastGetopt::NEscaping::Q;
+using NLastGetopt::NEscaping::QQ;
+using NLastGetopt::NEscaping::C;
+using NLastGetopt::NEscaping::CC;
+using NLastGetopt::NEscaping::S;
+using NLastGetopt::NEscaping::SS;
+using NLastGetopt::NEscaping::B;
+using NLastGetopt::NEscaping::BB;
+
+namespace NLastGetopt::NComp {
+#define L out.Line()
+#define I auto Y_GENERATE_UNIQUE_ID(indent) = out.Indent()
+
+    TCompleterManager::TCompleterManager(TStringBuf command)
+        : Command_(command)
+        , Id_(0)
+    {
+    }
+
+    TStringBuf TCompleterManager::GetCompleterID(const ICompleter* completer) {
+        return Queue_.emplace_back(TStringBuilder() << "_" << Command_ << "__completer_" << ++Id_, completer).first;
+    }
+
+    void TCompleterManager::GenerateZsh(TFormattedOutput& out) {
+        while (!Queue_.empty()) {
+            auto[name, completer] = Queue_.back();
+            Queue_.pop_back();
+
+            L << "(( $+functions[" << name << "] )) ||";
+            L << name << "() {";
+            {
+                I;
+                completer->GenerateZsh(out, *this);
+            }
+            L << "}";
+            L;
+        }
+    }
+
+    class TAlternativeCompleter: public ICompleter {
+    public:
+        TAlternativeCompleter(TVector<TAlternative>  alternatives)
+            : Alternatives_(std::move(alternatives))
+        {
+        }
+
+        void GenerateBash(TFormattedOutput& out) const override {
+            for (auto& alternative: Alternatives_) {
+                if (alternative.Completer != nullptr) {
+                    alternative.Completer->GenerateBash(out);
+                }
+            }
+        }
+
+        TStringBuf GenerateZshAction(TCompleterManager& manager) const override {
+            return manager.GetCompleterID(this);
+        }
+
+        void GenerateZsh(TFormattedOutput& out, TCompleterManager& manager) const override {
+            // We should use '_alternative' here, but it doesn't process escape codes in group descriptions,
+            // so we dispatch alternatives ourselves.
+
+            L << "local expl action";
+
+            size_t i = 0;
+            for (auto& alternative: Alternatives_) {
+                auto tag = "alt-" + ToString(++i);
+                auto action = alternative.Completer ? alternative.Completer->GenerateZshAction(manager) : TStringBuf();
+
+                L;
+
+                if (action.empty()) {
+                    L << "_message -e " << SS(tag) << " " << SS(alternative.Description);
+                } else if (action.StartsWith("((") && action.EndsWith("))")) {
+                    L << "action=" << action.substr(1, action.size() - 2);
+                    L << "_describe -t " << SS(tag) << " " << SS(alternative.Description) << " action -M 'r:|[_-]=* r:|=*'";
+                } else if (action.StartsWith("(") && action.EndsWith(")")) {
+                    L << "action=" << action << "";
+                    L << "_describe -t " << SS(tag) << " " << SS(alternative.Description) << " action -M 'r:|[_-]=* r:|=*'";
+                } else if (action.StartsWith(' ')) {
+                    L << action.substr(1);
+                } else {
+                    L << "_description " << SS(tag) << " expl " << SS(alternative.Description);
+                    TStringBuf word, args;
+                    action.Split(' ', word, args);
+                    L << word << " \"${expl[@]}\" " << args;
+                }
+            }
+        }
+
+    private:
+        TVector<TAlternative> Alternatives_;
+    };
+
+    ICompleterPtr Alternative(TVector<TAlternative> alternatives) {
+        return MakeSimpleShared<TAlternativeCompleter>(std::move(alternatives));
+    }
+
+    class TSimpleCompleter: public ICompleter {
+    public:
+        TSimpleCompleter(TString bashCode, TString action)
+            : BashCode(std::move(bashCode))
+            , Action(std::move(action))
+        {
+        }
+
+        void GenerateBash(TFormattedOutput& out) const override {
+            if (BashCode) {
+                L << BashCode;
+            }
+        }
+
+        TStringBuf GenerateZshAction(TCompleterManager&) const override {
+            return Action;
+        }
+
+        void GenerateZsh(TFormattedOutput&, TCompleterManager&) const override {
+            Y_FAIL("unreachable");
+        }
+
+    private:
+        TString BashCode;
+        TString Action;
+    };
+
+    ICompleterPtr Choice(TVector<TChoice> choices) {
+        auto bash = TStringBuilder() << "COMPREPLY+=( $(compgen -W '";
+        TStringBuf sep = "";
+        for (auto& choice : choices) {
+            bash << sep << B(choice.Choice);
+            sep = " ";
+        }
+        bash << "' -- ${cur}) )";
+
+        auto action = TStringBuilder();
+        action << "((";
+        for (auto& choice: choices) {
+            action << " " << SS(choice.Choice);
+            if (choice.Description) {{
+                action << ":" << SS(choice.Description);
+            }}
+        }
+        action << "))";
+        return MakeSimpleShared<TSimpleCompleter>(bash, action);
+    }
+
+    TString Compgen(TStringBuf flags) {
+        return TStringBuilder() << "COMPREPLY+=( $(compgen " << flags << " -- ${cur}) )";
+    }
+
+    ICompleterPtr Default() {
+        return MakeSimpleShared<TSimpleCompleter>("", "_default");
+    }
+
+    ICompleterPtr File(TString pattern) {
+        if (pattern) {
+            pattern = " -g " + SS(pattern);
+        }
+        return MakeSimpleShared<TSimpleCompleter>("", "_files" + pattern);
+    }
+
+    ICompleterPtr Directory() {
+        return MakeSimpleShared<TSimpleCompleter>("", "_files -/");
+    }
+
+    ICompleterPtr Host() {
+        return MakeSimpleShared<TSimpleCompleter>(Compgen("-A hostname"), "_hosts");
+    }
+
+    ICompleterPtr Pid() {
+        return MakeSimpleShared<TSimpleCompleter>("", "_pids");
+    }
+
+    ICompleterPtr User() {
+        return MakeSimpleShared<TSimpleCompleter>(Compgen("-A user"), "_users");
+    }
+
+    ICompleterPtr Group() {
+        // For some reason, OSX freezes when trying to perform completion for groups.
+        // You can try removing this ifdef and debugging it, but be prepared to force-shutdown your machine
+        // (and possibly reinstall OSX if force-shutdown breaks anything).
+#ifdef _darwin_
+        return MakeSimpleShared<TSimpleCompleter>("", "");
+#else
+        return MakeSimpleShared<TSimpleCompleter>(Compgen("-A group"), "_groups");
+#endif
+    }
+
+    ICompleterPtr Url() {
+        return MakeSimpleShared<TSimpleCompleter>("", "_urls");
+    }
+
+    ICompleterPtr Tty() {
+        return MakeSimpleShared<TSimpleCompleter>("", "_ttys");
+    }
+
+    ICompleterPtr NetInterface() {
+        return MakeSimpleShared<TSimpleCompleter>("", "_net_interfaces");
+    }
+
+    ICompleterPtr TimeZone() {
+        return MakeSimpleShared<TSimpleCompleter>("", "_time_zone");
+    }
+
+    ICompleterPtr Signal() {
+        return MakeSimpleShared<TSimpleCompleter>(Compgen("-A signal"), "_signals");
+    }
+
+    ICompleterPtr Domain() {
+        return MakeSimpleShared<TSimpleCompleter>("", "_domains");
+    }
+
+    namespace {
+        TCustomCompleter* Head = nullptr;
+        TStringBuf SpecialFlag = "---CUSTOM-COMPLETION---";
+    }
+
+    void TCustomCompleter::FireCustomCompleter(int argc, const char** argv) {
+        if (!argc) {
+            return;
+        }
+
+        for (int i = 1; i < argc - 4; ++i) {
+            if (SpecialFlag == argv[i]) {
+                auto name = TStringBuf(argv[i + 1]);
+                auto curIdx = FromString<int>(argv[i + 2]);
+                auto prefix = TStringBuf(argv[i + 3]);
+                auto suffix = TStringBuf(argv[i + 4]);
+
+                auto cur = TStringBuf();
+                if (0 <= curIdx && curIdx < i) {
+                    cur = TStringBuf(argv[curIdx]);
+                }
+                if (cur && !prefix && !suffix) {
+                    prefix = cur;  // bash does not send prefix and suffix
+                }
+
+                auto head = Head;
+                while (head) {
+                    if (head->GetUniqueName() == name) {
+                        head->GenerateCompletions(i, argv, curIdx, cur, prefix, suffix);
+                    }
+                    head = head->Next_;
+                }
+
+                exit(0);
+            }
+        }
+    }
+
+    void TCustomCompleter::RegisterCustomCompleter(TCustomCompleter* completer) noexcept {
+        Y_VERIFY(completer);
+        completer->Next_ = Head;
+        Head = completer;
+    }
+
+    void TCustomCompleter::AddCompletion(TStringBuf completion) {
+        Cout << completion << Endl;  // this was easy =)
+        // TODO: support option descriptions and messages
+    }
+
+    void TMultipartCustomCompleter::GenerateCompletions(int argc, const char** argv, int curIdx, TStringBuf cur, TStringBuf prefix, TStringBuf suffix) {
+        auto root = TStringBuf();
+        if (prefix.Contains(Sep_)) {
+            auto tmp = TStringBuf();
+            prefix.RSplit(Sep_, root, tmp);
+        }
+
+        if (root) {
+            Cout << root << Sep_ << Endl;
+        } else {
+            Cout << Endl;
+        }
+
+        Cout << Sep_ << Endl;
+
+        GenerateCompletionParts(argc, argv, curIdx, cur, prefix, suffix, root);
+    }
+
+    class TLaunchSelf: public ICompleter {
+    public:
+        TLaunchSelf(TCustomCompleter* completer)
+            : Completer_(completer)
+        {
+        }
+
+        void GenerateBash(TFormattedOutput& out) const override {
+            L << "IFS=$'\\n'";
+            L << "COMPREPLY+=( $(compgen -W \"$(${words[@]} " << SpecialFlag << " " << Completer_->GetUniqueName() << " \"${cword}\" \"\" \"\" 2> /dev/null)\" -- ${cur}) )";
+            L << "IFS=$' \\t\\n'";
+        }
+
+        TStringBuf GenerateZshAction(TCompleterManager& manager) const override {
+            return manager.GetCompleterID(this);
+        }
+
+        void GenerateZsh(TFormattedOutput& out, TCompleterManager&) const override {
+            L << "compadd ${@} ${expl[@]} -- \"${(@f)$(${words_orig[@]} " << SpecialFlag << " " << Completer_->GetUniqueName() << " \"${current_orig}\" \"${prefix_orig}\" \"${suffix_orig}\" 2> /dev/null)}\"";
+        }
+
+    private:
+        TCustomCompleter* Completer_;
+    };
+
+    ICompleterPtr LaunchSelf(TCustomCompleter& completer) {
+        return MakeSimpleShared<TLaunchSelf>(&completer);
+    }
+
+    class TLaunchSelfMultiPart: public ICompleter {
+    public:
+        TLaunchSelfMultiPart(TCustomCompleter* completer)
+            : Completer_(completer)
+        {
+        }
+
+        void GenerateBash(TFormattedOutput& out) const override {
+            L << "IFS=$'\\n'";
+            L << "items=( $(${words[@]} " << SpecialFlag << " " << Completer_->GetUniqueName() << " \"${cword}\" \"\" \"\" 2> /dev/null) )";
+            L << "candidates=$(compgen -W \"${items[*]:1}\" -- \"$cur\")";
+            L << "COMPREPLY+=( $candidates )";
+            L << "[[ $candidates == *\"${items[1]}\" ]] && need_space=\"\"";
+            L << "IFS=$' \\t\\n'";
+        }
+
+        TStringBuf GenerateZshAction(TCompleterManager& manager) const override {
+            return manager.GetCompleterID(this);
+        }
+
+        void GenerateZsh(TFormattedOutput& out, TCompleterManager&) const override {
+            L << "local items=( \"${(@f)$(${words_orig[@]} " << SpecialFlag << " " << Completer_->GetUniqueName() << " \"${current_orig}\" \"${prefix_orig}\" \"${suffix_orig}\" 2> /dev/null)}\" )";
+            L;
+            L << "local rempat=${items[1]}";
+            L << "shift items";
+            L;
+            L << "local sep=${items[1]}";
+            L << "shift items";
+            L;
+            L << "local files=( ${items:#*\"${sep}\"} )";
+            L << "local filenames=( ${files#\"${rempat}\"} )";
+            L << "local dirs=( ${(M)items:#*\"${sep}\"} )";
+            L << "local dirnames=( ${dirs#\"${rempat}\"} )";
+            L;
+            L << "local need_suf";
+            L << "compset -S \"${sep}*\" || need_suf=\"1\"";
+            L;
+            L << "compadd ${@} ${expl[@]} -d filenames -- ${(q)files}";
+            L << "compadd ${@} ${expl[@]} ${need_suf:+-S\"${sep}\"} -q -d dirnames -- ${(q)dirs%\"${sep}\"}";
+        }
+
+    private:
+        TCustomCompleter* Completer_;
+    };
+
+    ICompleterPtr LaunchSelfMultiPart(TCustomCompleter& completer) {
+        return MakeSimpleShared<TLaunchSelfMultiPart>(&completer);
+    }
+
+#undef I
+#undef L
+}