Merge remote-tracking branch 'upstream/master' into path-info

1 files changed, 119 insertions, 20 deletions

diff --git a/src/libutil/args.cc b/src/libutil/args.cc
index 1f0e4f803..986c5d1cd 100644
--- a/src/libutil/args.cc
+++ b/src/libutil/args.cc
@@ -1,6 +1,8 @@
 #include "args.hh"
 #include "hash.hh"
 
+#include <glob.h>
+
 namespace nix {
 
 void Args::addFlag(Flag && flag_)
@@ -13,6 +15,20 @@ void Args::addFlag(Flag && flag_)
     if (flag->shortName) shortFlags[flag->shortName] = flag;
 }
 
+bool pathCompletions = false;
+std::shared_ptr<std::set<std::string>> completions;
+
+std::string completionMarker = "___COMPLETE___";
+
+std::optional<std::string> needsCompletion(std::string_view s)
+{
+    if (!completions) return {};
+    auto i = s.find(completionMarker);
+    if (i != std::string::npos)
+        return std::string(s.begin(), i);
+    return {};
+}
+
 void Args::parseCmdline(const Strings & _cmdline)
 {
     Strings pendingArgs;
@@ -20,6 +36,14 @@ void Args::parseCmdline(const Strings & _cmdline)
 
     Strings cmdline(_cmdline);
 
+    if (auto s = getEnv("NIX_GET_COMPLETIONS")) {
+        size_t n = std::stoi(*s);
+        assert(n > 0 && n <= cmdline.size());
+        *std::next(cmdline.begin(), n - 1) += completionMarker;
+        completions = std::make_shared<decltype(completions)::element_type>();
+        verbosity = lvlError;
+    }
+
     for (auto pos = cmdline.begin(); pos != cmdline.end(); ) {
 
         auto arg = *pos;
@@ -63,7 +87,7 @@ void Args::printHelp(const string & programName, std::ostream & out)
     for (auto & exp : expectedArgs) {
         std::cout << renderLabels({exp.label});
         // FIXME: handle arity > 1
-        if (exp.arity == 0) std::cout << "...";
+        if (exp.handler.arity == ArityAny) std::cout << "...";
         if (exp.optional) std::cout << "?";
     }
     std::cout << "\n";
@@ -99,18 +123,32 @@ bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
     auto process = [&](const std::string & name, const Flag & flag) -> bool {
         ++pos;
         std::vector<std::string> args;
+        bool anyCompleted = false;
         for (size_t n = 0 ; n < flag.handler.arity; ++n) {
             if (pos == end) {
                 if (flag.handler.arity == ArityAny) break;
                 throw UsageError("flag '%s' requires %d argument(s)", name, flag.handler.arity);
             }
+            if (flag.completer)
+                if (auto prefix = needsCompletion(*pos)) {
+                    anyCompleted = true;
+                    flag.completer(n, *prefix);
+                }
             args.push_back(*pos++);
         }
-        flag.handler.fun(std::move(args));
+        if (!anyCompleted)
+            flag.handler.fun(std::move(args));
         return true;
     };
 
     if (string(*pos, 0, 2) == "--") {
+        if (auto prefix = needsCompletion(*pos)) {
+            for (auto & [name, flag] : longFlags) {
+                if (!hiddenCategories.count(flag->category)
+                    && hasPrefix(name, std::string(*prefix, 2)))
+                    completions->insert("--" + name);
+            }
+        }
         auto i = longFlags.find(string(*pos, 2));
         if (i == longFlags.end()) return false;
         return process("--" + i->first, *i->second);
@@ -123,6 +161,14 @@ bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
         return process(std::string("-") + c, *i->second);
     }
 
+    if (auto prefix = needsCompletion(*pos)) {
+        if (prefix == "-") {
+            completions->insert("--");
+            for (auto & [flag, _] : shortFlags)
+                completions->insert(std::string("-") + flag);
+        }
+    }
+
     return false;
 }
 
@@ -138,12 +184,17 @@ bool Args::processArgs(const Strings & args, bool finish)
 
     bool res = false;
 
-    if ((exp.arity == 0 && finish) ||
-        (exp.arity > 0 && args.size() == exp.arity))
+    if ((exp.handler.arity == ArityAny && finish) ||
+        (exp.handler.arity != ArityAny && args.size() == exp.handler.arity))
     {
         std::vector<std::string> ss;
-        for (auto & s : args) ss.push_back(s);
-        exp.handler(std::move(ss));
+        for (const auto &[n, s] : enumerate(args)) {
+            ss.push_back(s);
+            if (exp.completer)
+                if (auto prefix = needsCompletion(s))
+                    exp.completer(n, *prefix);
+        }
+        exp.handler.fun(ss);
         expectedArgs.pop_front();
         res = true;
     }
@@ -154,6 +205,13 @@ bool Args::processArgs(const Strings & args, bool finish)
     return res;
 }
 
+static void hashTypeCompleter(size_t index, std::string_view prefix) 
+{
+    for (auto & type : hashTypes)
+        if (hasPrefix(type, prefix))
+            completions->insert(type);
+}
+
 Args::Flag Args::Flag::mkHashTypeFlag(std::string && longName, HashType * ht)
 {
     return Flag {
@@ -162,7 +220,8 @@ Args::Flag Args::Flag::mkHashTypeFlag(std::string && longName, HashType * ht)
         .labels = {"hash-algo"},
         .handler = {[ht](std::string s) {
             *ht = parseHashType(s);
-        }}
+        }},
+        .completer = hashTypeCompleter
     };
 }
 
@@ -174,10 +233,42 @@ Args::Flag Args::Flag::mkHashTypeOptFlag(std::string && longName, std::optional<
         .labels = {"hash-algo"},
         .handler = {[oht](std::string s) {
             *oht = std::optional<HashType> { parseHashType(s) };
-        }}
+        }},
+        .completer = hashTypeCompleter
     };
 }
 
+static void completePath(std::string_view prefix, bool onlyDirs)
+{
+    pathCompletions = true;
+    glob_t globbuf;
+    int flags = GLOB_NOESCAPE | GLOB_TILDE;
+    #ifdef GLOB_ONLYDIR
+    if (onlyDirs)
+        flags |= GLOB_ONLYDIR;
+    #endif
+    if (glob((std::string(prefix) + "*").c_str(), flags, nullptr, &globbuf) == 0) {
+        for (size_t i = 0; i < globbuf.gl_pathc; ++i) {
+            if (onlyDirs) {
+                auto st = lstat(globbuf.gl_pathv[i]);
+                if (!S_ISDIR(st.st_mode)) continue;
+            }
+            completions->insert(globbuf.gl_pathv[i]);
+        }
+        globfree(&globbuf);
+    }
+}
+
+void completePath(size_t, std::string_view prefix)
+{
+    completePath(prefix, false);
+}
+
+void completeDir(size_t, std::string_view prefix)
+{
+    completePath(prefix, true);
+}
+
 Strings argvToStrings(int argc, char * * argv)
 {
     Strings args;
@@ -225,18 +316,26 @@ void Command::printHelp(const string & programName, std::ostream & out)
 MultiCommand::MultiCommand(const Commands & commands)
     : commands(commands)
 {
-    expectedArgs.push_back(ExpectedArg{"command", 1, true, [=](std::vector<std::string> ss) {
-        assert(!command);
-        auto cmd = ss[0];
-        if (auto alias = get(deprecatedAliases, cmd)) {
-            warn("'%s' is a deprecated alias for '%s'", cmd, *alias);
-            cmd = *alias;
-        }
-        auto i = commands.find(cmd);
-        if (i == commands.end())
-            throw UsageError("'%s' is not a recognised command", cmd);
-        command = {cmd, i->second()};
-    }});
+    expectArgs({
+        .label = "command",
+        .optional = true,
+        .handler = {[=](std::string s) {
+            assert(!command);
+            if (auto alias = get(deprecatedAliases, s)) {
+                warn("'%s' is a deprecated alias for '%s'", s, *alias);
+                s = *alias;
+            }
+            if (auto prefix = needsCompletion(s)) {
+                for (auto & [name, command] : commands)
+                    if (hasPrefix(name, *prefix))
+                        completions->insert(name);
+            }
+            auto i = commands.find(s);
+            if (i == commands.end())
+                throw UsageError("'%s' is not a recognised command", s);
+            command = {s, i->second()};
+        }}
+    });
 
     categories[Command::catDefault] = "Available commands";
 }