Merge pull request #4128 from tweag/extended-completions

Add a zsh completion script

2 files changed, 31 insertions, 10 deletions

diff --git a/src/libutil/args.cc b/src/libutil/args.cc
index 2760b830b..66d8df085 100644
--- a/src/libutil/args.cc
+++ b/src/libutil/args.cc
@@ -17,8 +17,19 @@ void Args::addFlag(Flag && flag_)
     if (flag->shortName) shortFlags[flag->shortName] = flag;
 }
 
+void Completions::add(std::string completion, std::string description)
+{
+    insert(Completion{
+            .completion = completion,
+            .description = description
+            });
+}
+
+bool Completion::operator<(const Completion & other) const
+{ return completion < other.completion || (completion == other.completion && description < other.description); }
+
 bool pathCompletions = false;
-std::shared_ptr<std::set<std::string>> completions;
+std::shared_ptr<Completions> completions;
 
 std::string completionMarker = "___COMPLETE___";
 
@@ -148,7 +159,7 @@ bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
             for (auto & [name, flag] : longFlags) {
                 if (!hiddenCategories.count(flag->category)
                     && hasPrefix(name, std::string(*prefix, 2)))
-                    completions->insert("--" + name);
+                    completions->add("--" + name, flag->description);
             }
         }
         auto i = longFlags.find(string(*pos, 2));
@@ -165,9 +176,9 @@ bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
 
     if (auto prefix = needsCompletion(*pos)) {
         if (prefix == "-") {
-            completions->insert("--");
-            for (auto & [flag, _] : shortFlags)
-                completions->insert(std::string("-") + flag);
+            completions->add("--");
+            for (auto & [flagName, flag] : shortFlags)
+                completions->add(std::string("-") + flagName, flag->description);
         }
     }
 
@@ -244,11 +255,11 @@ nlohmann::json Args::toJSON()
     return res;
 }
 
-static void hashTypeCompleter(size_t index, std::string_view prefix) 
+static void hashTypeCompleter(size_t index, std::string_view prefix)
 {
     for (auto & type : hashTypes)
         if (hasPrefix(type, prefix))
-            completions->insert(type);
+            completions->add(type);
 }
 
 Args::Flag Args::Flag::mkHashTypeFlag(std::string && longName, HashType * ht)
@@ -292,7 +303,7 @@ static void _completePath(std::string_view prefix, bool onlyDirs)
                 auto st = lstat(globbuf.gl_pathv[i]);
                 if (!S_ISDIR(st.st_mode)) continue;
             }
-            completions->insert(globbuf.gl_pathv[i]);
+            completions->add(globbuf.gl_pathv[i]);
         }
         globfree(&globbuf);
     }
@@ -385,7 +396,7 @@ MultiCommand::MultiCommand(const Commands & commands)
             if (auto prefix = needsCompletion(s)) {
                 for (auto & [name, command] : commands)
                     if (hasPrefix(name, *prefix))
-                        completions->insert(name);
+                        completions->add(name);
             }
             auto i = commands.find(s);
             if (i == commands.end())
diff --git a/src/libutil/args.hh b/src/libutil/args.hh
index f41242e17..26f1bc11b 100644
--- a/src/libutil/args.hh
+++ b/src/libutil/args.hh
@@ -283,7 +283,17 @@ typedef std::vector<std::pair<std::string, std::string>> Table2;
 
 void printTable(std::ostream & out, const Table2 & table);
 
-extern std::shared_ptr<std::set<std::string>> completions;
+struct Completion {
+    std::string completion;
+    std::string description;
+
+    bool operator<(const Completion & other) const;
+};
+class Completions : public std::set<Completion> {
+public:
+    void add(std::string completion, std::string description = "");
+};
+extern std::shared_ptr<Completions> completions;
 extern bool pathCompletions;
 
 std::optional<std::string> needsCompletion(std::string_view s);