feat: implement shell completion helpers

crates/rattler_shell/src/lib.rs

@@ -1,6 +1,5 @@
-#![deny(missing_docs)]
-
 //! This crate provides helper functions to activate and deactivate virtual environments.
+#![deny(missing_docs)]
 
 pub mod activation;
 pub mod run;

crates/rattler_shell/src/shell/mod.rs

@@ -51,6 +51,14 @@ pub trait Shell {
     /// Run a script in the current shell.
     fn run_script(&self, f: &mut impl Write, path: &Path) -> std::fmt::Result;
 
+    /// Source completion scripts for the shell from a given prefix path.
+    /// Note: the `completions_dir` is the directory where the completions are stored.
+    /// You can use [`Self::completion_script_location`] to get the correct location for a given
+    /// shell type.
+    fn source_completions(&self, _f: &mut impl Write, _completions_dir: &Path) -> std::fmt::Result {
+        Ok(())
+    }
+
     /// Test to see if the path can be executed by the shell, based on the
     /// extension of the path.
     fn can_run_script(&self, path: &Path) -> bool {
@@ -162,6 +170,17 @@ pub trait Shell {
     fn line_ending(&self) -> &str {
         "\n"
     }
+
+    /// Return the location where completion scripts are found in a Conda environment.
+    ///
+    /// - bash: `share/bash-completion/completions`
+    /// - zsh: `share/zsh/site-functions`
+    /// - fish: `share/fish/vendor_completions.d`
+    ///
+    /// The return value must be joined with `prefix.join(completion_script_location())`.
+    fn completion_script_location(&self) -> Option<&'static Path> {
+        None
+    }
 }
 
 /// Convert a native PATH on Windows to a Unix style path using cygpath.
@@ -258,6 +277,18 @@ impl Shell for Bash {
         "sh"
     }
 
+    fn completion_script_location(&self) -> Option<&'static Path> {
+        Some(Path::new("share/bash-completion/completions"))
+    }
+
+    fn source_completions(&self, f: &mut impl Write, completions_dir: &Path) -> std::fmt::Result {
+        if completions_dir.exists() {
+            let completions_glob = completions_dir.join("*");
+            writeln!(f, "source {}", completions_glob.to_string_lossy())?;
+        }
+        Ok(())
+    }
+
     fn executable(&self) -> &str {
         "bash"
     }
@@ -307,6 +338,19 @@ impl Shell for Zsh {
         cmd.arg(path);
         cmd
     }
+
+    fn completion_script_location(&self) -> Option<&'static Path> {
+        Some(Path::new("share/zsh/site-functions"))
+    }
+
+    fn source_completions(&self, f: &mut impl Write, completions_dir: &Path) -> std::fmt::Result {
+        if completions_dir.exists() {
+            writeln!(f, "fpath+=({})", completions_dir.to_string_lossy())?;
+            writeln!(f, "autoload -Uz compinit")?;
+            writeln!(f, "compinit")?;
+        }
+        Ok(())
+    }
 }
 
 /// A [`Shell`] implementation for the Xonsh shell.
@@ -352,6 +396,10 @@ impl Shell for Xonsh {
         cmd.arg(path);
         cmd
     }
+
+    fn completion_script_location(&self) -> Option<&'static Path> {
+        None
+    }
 }
 
 /// A [`Shell`] implementation for the cmd.exe shell.
@@ -527,6 +575,21 @@ impl Shell for Fish {
         cmd.arg(path);
         cmd
     }
+
+    fn completion_script_location(&self) -> Option<&'static Path> {
+        Some(Path::new("share/fish/vendor_completions.d"))
+    }
+
+    fn source_completions(&self, f: &mut impl Write, completions_dir: &Path) -> std::fmt::Result {
+        if completions_dir.exists() {
+            // glob all files in the completions directory using fish
+            let completions_glob = completions_dir.join("*");
+            writeln!(f, "for file in {}", completions_glob.to_string_lossy())?;
+            writeln!(f, "    source $file")?;
+            writeln!(f, "end")?;
+        }
+        Ok(())
+    }
 }
 
 fn escape_backslashes(s: &str) -> String {
@@ -603,6 +666,10 @@ impl Shell for NuShell {
         cmd.arg(path);
         cmd
     }
+
+    fn completion_script_location(&self) -> Option<&'static Path> {
+        None
+    }
 }
 
 /// A generic [`Shell`] implementation for concrete shell types.
@@ -810,6 +877,16 @@ impl<T: Shell + 'static> ShellScript<T> {
         Ok(self)
     }
 
+    /// Source completion scripts for the shell from a given directory with completion scripts.
+    pub fn source_completions(
+        &mut self,
+        completions_dir: &Path,
+    ) -> Result<&mut Self, std::fmt::Error> {
+        self.shell
+            .source_completions(&mut self.contents, completions_dir)?;
+        Ok(self)
+    }
+
     /// Add contents to the script. The contents will be added as is, so make
     /// sure to format it correctly for the shell.
     pub fn append_script(&mut self, script: &Self) -> &mut Self {