feat: multiple commands for shell mode

src/main.rs

@@ -78,7 +78,7 @@ fn index_database(command: &str, picker: &str) -> Option<String> {
 
 fn run_command_or_open_shell(
     use_channel: bool,
-    choice: &str,
+    choice: Vec<String>,
     command: &str,
     trail: &[String],
     nixpkgs_flake: &str,
@@ -92,9 +92,12 @@ fn run_command_or_open_shell(
     ]);
 
     if use_channel {
-        run_cmd.args(["-f", "<nixpkgs>", choice]);
+        run_cmd.args(["-f", "<nixpkgs>"]);
+        run_cmd.args(choice);
     } else {
-        run_cmd.args([format!("{nixpkgs_flake}#{choice}")]);
+        for prog in choice {
+            run_cmd.args([format!("{nixpkgs_flake}#{prog}")]);
+        }
     }
 
     if !command.is_empty() {
@@ -110,22 +113,36 @@ fn run_command_or_open_shell(
 fn main() -> ExitCode {
     let args = Opt::parse();
 
-    let mut cache = Cache::new();
-    if let Err(ref e) = cache {
-        eprintln!("failed to initialize cache, disabling related functionality: {e}");
-    }
+    let mut cache = Cache::new()
+        .map_err(|e| {
+            eprintln!("failed to initialize cache, disabling related functionality: {e}");
+        })
+        .ok();
 
     if args.update {
         eprintln!("\"comma --update\" has been deprecated. either obtain a prebuilt database from https://github.com/Mic92/nix-index-database or use \"nix run 'nixpkgs#nix-index' --extra-experimental-features 'nix-command flakes'\"");
         index::update_database();
     }
 
     if args.empty_cache {
-        if let Ok(ref mut cache) = cache {
+        if let Some(ref mut cache) = cache {
             cache.empty();
         }
     }
 
+    let use_channel = env::var("NIX_PATH").is_ok_and(|p| p.contains("nixpkgs="));
+
+    if !args.shell.is_empty() {
+        return open_shell(
+            args.shell,
+            args.delete_entry,
+            &args.picker,
+            cache,
+            use_channel,
+            args.nixpkgs_flake,
+        );
+    }
+
     // The command may not be given if `--update` was specified.
     if args.cmd.is_empty() {
         return if args.update || args.empty_cache {
@@ -138,35 +155,14 @@ fn main() -> ExitCode {
     let command = &args.cmd[0];
     let trail = &args.cmd[1..];
 
-    if args.delete_entry {
-        if let Ok(ref mut cache) = cache {
-            cache.delete(command);
-        }
-    }
-
-    let derivation = match cache {
-        Ok(mut cache) => cache.query(command).or_else(|| {
-            index_database(command, &args.picker).map(|derivation| {
-                cache.update(command, &derivation);
-                derivation
-            })
-        }),
-        Err(_) => index_database(command, &args.picker),
-    };
-
-    let derivation = match derivation {
-        Some(d) => d,
-        None => return ExitCode::FAILURE,
+    let Some(derivation) = find_program(command, &args.picker, cache.as_mut(), args.delete_entry)
+    else {
+        eprint!("Error");
+        return ExitCode::FAILURE;
     };
-
+    // Explicitly drop cache because exec doesn't call destructors
+    drop(cache);
     let basename = derivation.rsplit('.').last().unwrap();
-
-    let use_channel = match env::var("NIX_PATH") {
-        Ok(val) => val,
-        Err(_) => String::new(),
-    }
-    .contains("nixpkgs=");
-
     if args.print_package {
         println!(
             "Package that contains executable /bin/{}: {}",
@@ -178,19 +174,10 @@ fn main() -> ExitCode {
         Command::new("nix-env")
             .args(["-f", "<nixpkgs>", "-iA", basename])
             .exec();
-    } else if args.shell {
-        let shell_cmd = shell::select_shell_from_pid(process::id()).unwrap_or("bash".into());
-        run_command_or_open_shell(
-            use_channel,
-            &derivation,
-            &shell_cmd,
-            &[],
-            &args.nixpkgs_flake,
-        );
     } else if args.print_path {
         run_command_or_open_shell(
             use_channel,
-            &derivation,
+            vec![derivation],
             "sh",
             &[
                 String::from("-c"),
@@ -201,7 +188,7 @@ fn main() -> ExitCode {
     } else {
         run_command_or_open_shell(
             use_channel,
-            &derivation,
+            vec![derivation],
             command,
             trail,
             &args.nixpkgs_flake,
@@ -211,6 +198,60 @@ fn main() -> ExitCode {
     ExitCode::SUCCESS
 }
 
+fn open_shell(
+    commands: Vec<String>,
+    empty_cache: bool,
+    picker: &str,
+    mut cache: Option<Cache>,
+    use_channel: bool,
+    nixpkgs_flake: String,
+) -> ExitCode {
+    let shell_cmd = shell::select_shell_from_pid(process::id()).unwrap_or("bash".into());
+    let mut programs = Vec::new();
+    for prog in commands {
+        let Some(derivation) = find_program(&prog, picker, cache.as_mut(), empty_cache) else {
+            eprint!("Couldn't find program for {prog}");
+            return ExitCode::FAILURE;
+        };
+        programs.push(derivation);
+    }
+
+    // Explicitly drop cache because exec doesn't call destructors
+    drop(cache);
+
+    run_command_or_open_shell(
+        use_channel,
+        programs,
+        &shell_cmd,
+        &[],
+        nixpkgs_flake.as_str(),
+    );
+    ExitCode::SUCCESS
+}
+
+/// Try to find and select a program from a command
+fn find_program(
+    command: &str,
+    picker: &str,
+    cache: Option<&mut Cache>,
+    delete_entry: bool,
+) -> Option<String> {
+    cache.map_or_else(
+        || index_database(command, picker),
+        |cache| {
+            if delete_entry {
+                cache.delete(command);
+            }
+            cache.query(command).or_else(|| {
+                index_database(command, picker).map(|derivation| {
+                    cache.update(command, &derivation);
+                    derivation
+                })
+            })
+        },
+    )
+}
+
 /// Runs programs without installing them
 #[derive(Parser)]
 #[clap(version = crate_version!(), trailing_var_arg = true)]
@@ -220,8 +261,8 @@ struct Opt {
     install: bool,
 
     /// Open a shell containing the derivation containing the executable
-    #[clap(short, long)]
-    shell: bool,
+    #[clap(short, long, conflicts_with_all(["cmd", "install"]), num_args(1..))]
+    shell: Vec<String>,
 
     #[clap(short = 'P', long, env = "COMMA_PICKER", default_value = "fzy")]
     picker: String,
@@ -256,6 +297,6 @@ struct Opt {
     delete_entry: bool,
 
     /// Command to run
-    #[clap(required_unless_present_any = ["update", "empty_cache"], name = "cmd")]
+    #[clap(required_unless_present_any = ["update", "empty_cache", "shell"], name = "cmd")]
     cmd: Vec<String>,
 }