Merge pull request #6693 from ncfavier/complete-flake-inputs

Improve shell completion of flake inputs

6 files changed, 57 insertions, 29 deletions

diff --git a/src/libcmd/command.hh b/src/libcmd/command.hh
index 6bd38810b..3b4b40981 100644
--- a/src/libcmd/command.hh
+++ b/src/libcmd/command.hh
@@ -78,10 +78,16 @@ struct MixFlakeOptions : virtual Args, EvalCommand
 {
     flake::LockFlags lockFlags;
 
+    std::optional<std::string> needsFlakeInputCompletion = {};
+
     MixFlakeOptions();
 
-    virtual std::optional<FlakeRef> getFlakeRefForCompletion()
+    virtual std::vector<std::string> getFlakesForCompletion()
     { return {}; }
+
+    void completeFlakeInput(std::string_view prefix);
+
+    void completionHook() override;
 };
 
 struct SourceExprCommand : virtual Args, MixFlakeOptions
@@ -121,7 +127,7 @@ struct InstallablesCommand : virtual Args, SourceExprCommand
 
     virtual bool useDefaultInstallables() { return true; }
 
-    std::optional<FlakeRef> getFlakeRefForCompletion() override;
+    std::vector<std::string> getFlakesForCompletion() override;
 
 protected:
 
@@ -137,9 +143,9 @@ struct InstallableCommand : virtual Args, SourceExprCommand
 
     void prepare() override;
 
-    std::optional<FlakeRef> getFlakeRefForCompletion() override
+    std::vector<std::string> getFlakesForCompletion() override
     {
-        return parseFlakeRefWithFragment(_installable, absPath(".")).first;
+        return {_installable};
     }
 
 private:
diff --git a/src/libcmd/installables.cc b/src/libcmd/installables.cc
index 03c8f82b3..59162c4df 100644
--- a/src/libcmd/installables.cc
+++ b/src/libcmd/installables.cc
@@ -23,17 +23,6 @@
 
 namespace nix {
 
-void completeFlakeInputPath(
-    ref<EvalState> evalState,
-    const FlakeRef & flakeRef,
-    std::string_view prefix)
-{
-    auto flake = flake::getFlake(*evalState, flakeRef, true);
-    for (auto & input : flake.inputs)
-        if (hasPrefix(input.first, prefix))
-            completions->add(input.first);
-}
-
 MixFlakeOptions::MixFlakeOptions()
 {
     auto category = "Common flake-related options";
@@ -86,8 +75,7 @@ MixFlakeOptions::MixFlakeOptions()
             lockFlags.inputUpdates.insert(flake::parseInputPath(s));
         }},
         .completer = {[&](size_t, std::string_view prefix) {
-            if (auto flakeRef = getFlakeRefForCompletion())
-                completeFlakeInputPath(getEvalState(), *flakeRef, prefix);
+            needsFlakeInputCompletion = {std::string(prefix)};
         }}
     });
 
@@ -103,12 +91,10 @@ MixFlakeOptions::MixFlakeOptions()
                 parseFlakeRef(flakeRef, absPath("."), true));
         }},
         .completer = {[&](size_t n, std::string_view prefix) {
-            if (n == 0) {
-                if (auto flakeRef = getFlakeRefForCompletion())
-                    completeFlakeInputPath(getEvalState(), *flakeRef, prefix);
-            } else if (n == 1) {
+            if (n == 0)
+                needsFlakeInputCompletion = {std::string(prefix)};
+            else if (n == 1)
                 completeFlakeRef(getEvalState()->store, prefix);
-            }
         }}
     });
 
@@ -139,6 +125,24 @@ MixFlakeOptions::MixFlakeOptions()
     });
 }
 
+void MixFlakeOptions::completeFlakeInput(std::string_view prefix)
+{
+    auto evalState = getEvalState();
+    for (auto & flakeRefS : getFlakesForCompletion()) {
+        auto flakeRef = parseFlakeRefWithFragment(expandTilde(flakeRefS), absPath(".")).first;
+        auto flake = flake::getFlake(*evalState, flakeRef, true);
+        for (auto & input : flake.inputs)
+            if (hasPrefix(input.first, prefix))
+                completions->add(input.first);
+    }
+}
+
+void MixFlakeOptions::completionHook()
+{
+    if (auto & prefix = needsFlakeInputCompletion)
+        completeFlakeInput(*prefix);
+}
+
 SourceExprCommand::SourceExprCommand(bool supportReadOnlyMode)
 {
     addFlag({
@@ -1048,14 +1052,14 @@ Installables InstallablesCommand::load() {
     return parseInstallables(getStore(), _installables);
 }
 
-std::optional<FlakeRef> InstallablesCommand::getFlakeRefForCompletion()
+std::vector<std::string> InstallablesCommand::getFlakesForCompletion()
 {
     if (_installables.empty()) {
         if (useDefaultInstallables())
-            return parseFlakeRefWithFragment(".", absPath(".")).first;
+            return {"."};
         return {};
     }
-    return parseFlakeRefWithFragment(_installables.front(), absPath(".")).first;
+    return _installables;
 }
 
 InstallableCommand::InstallableCommand(bool supportReadOnlyMode)
diff --git a/src/libutil/args.cc b/src/libutil/args.cc
index 4b8c55686..44b63f0f6 100644
--- a/src/libutil/args.cc
+++ b/src/libutil/args.cc
@@ -124,7 +124,7 @@ bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
         bool anyCompleted = false;
         for (size_t n = 0 ; n < flag.handler.arity; ++n) {
             if (pos == end) {
-                if (flag.handler.arity == ArityAny) break;
+                if (flag.handler.arity == ArityAny || anyCompleted) break;
                 throw UsageError("flag '%s' requires %d argument(s)", name, flag.handler.arity);
             }
             if (auto prefix = needsCompletion(*pos)) {
@@ -362,6 +362,14 @@ bool MultiCommand::processArgs(const Strings & args, bool finish)
         return Args::processArgs(args, finish);
 }
 
+void MultiCommand::completionHook()
+{
+    if (command)
+        return command->second->completionHook();
+    else
+        return Args::completionHook();
+}
+
 nlohmann::json MultiCommand::toJSON()
 {
     auto cmds = nlohmann::json::object();
diff --git a/src/libutil/args.hh b/src/libutil/args.hh
index 07c017719..84866f12b 100644
--- a/src/libutil/args.hh
+++ b/src/libutil/args.hh
@@ -148,6 +148,11 @@ protected:
        argument (if any) have been processed. */
     virtual void initialFlagsProcessed() {}
 
+    /* Called after the command line has been processed if we need to generate
+       completions. Useful for commands that need to know the whole command line
+       in order to know what completions to generate. */
+    virtual void completionHook() { }
+
 public:
 
     void addFlag(Flag && flag);
@@ -223,6 +228,8 @@ public:
 
     bool processArgs(const Strings & args, bool finish) override;
 
+    void completionHook() override;
+
     nlohmann::json toJSON() override;
 };
 
diff --git a/src/nix/flake.cc b/src/nix/flake.cc
index c66b24e5c..e01bc6d10 100644
--- a/src/nix/flake.cc
+++ b/src/nix/flake.cc
@@ -50,9 +50,9 @@ public:
         return flake::lockFlake(*getEvalState(), getFlakeRef(), lockFlags);
     }
 
-    std::optional<FlakeRef> getFlakeRefForCompletion() override
+    std::vector<std::string> getFlakesForCompletion() override
     {
-        return getFlakeRef();
+        return {flakeUrl};
     }
 };
 
diff --git a/src/nix/main.cc b/src/nix/main.cc
index 17c92ebc6..a8404a2ea 100644
--- a/src/nix/main.cc
+++ b/src/nix/main.cc
@@ -347,7 +347,10 @@ void mainWrapped(int argc, char * * argv)
         if (!completions) throw;
     }
 
-    if (completions) return;
+    if (completions) {
+        args.completionHook();
+        return;
+    }
 
     if (args.showVersion) {
         printVersion(programName);