command-not-found: make NIX_AUTO_RUN work when multiple choices

1 files changed, 25 insertions, 4 deletions

diff --git a/nixos/modules/programs/command-not-found/command-not-found.pl b/nixos/modules/programs/command-not-found/command-not-found.pl
index 220d057b7f4f..5f4b30d90190 100644
--- a/nixos/modules/programs/command-not-found/command-not-found.pl
+++ b/nixos/modules/programs/command-not-found/command-not-found.pl
@@ -21,9 +21,11 @@ my $res = $dbh->selectall_arrayref(
     "select package from Programs where system = ? and name = ?",
     { Slice => {} }, $system, $program);
 
-if (!defined $res || scalar @$res == 0) {
+my $len = !defined $res ? 0 : scalar @$res;
+
+if ($len == 0) {
     print STDERR "$program: command not found\n";
-} elsif (scalar @$res == 1) {
+} elsif ($len == 1) {
     my $package = @$res[0]->{package};
     if ($ENV{"NIX_AUTO_RUN"} // "") {
         exec("nix-shell", "-p", $package, "--run", shell_quote("exec", @ARGV));
@@ -35,11 +37,30 @@ ephemeral shell by typing:
 EOF
     }
 } else {
-    print STDERR <<EOF;
+    if ($ENV{"NIX_AUTO_RUN"} // "") {
+        print STDERR "Select a package that provides '$program':\n";
+        for my $i (0 .. $len - 1) {
+            print STDERR "  [", $i + 1, "]: @$res[$i]->{package}\n";
+        }
+        my $choice = 0;
+        while (1) { # exec will break this loop
+            no warnings "numeric";
+            print STDERR "Your choice [1-${len}]: ";
+            # 0 can be invalid user input like non-number string
+            # so we start from 1
+            $choice = <STDIN> + 0;
+            if (1 <= $choice && $choice <= $len) {
+                exec("nix-shell", "-p", @$res[$choice - 1]->{package},
+                    "--run", shell_quote("exec", @ARGV));
+            }
+        }
+    } else {
+        print STDERR <<EOF;
 The program '$program' is not in your PATH. It is provided by several packages.
 You can make it available in an ephemeral shell by typing one of the following:
 EOF
-    print STDERR "  nix-shell -p $_->{package}\n" foreach @$res;
+        print STDERR "  nix-shell -p $_->{package}\n" foreach @$res;
+    }
 }
 
 exit 127;