allow specifying #! inside shell

fixes #51

src/script.rs

@@ -1,4 +1,4 @@
-use std::fs::read_to_string;
+use std::{fs::read_to_string, iter::once};
 
 use bollard::Docker;
 use camino::Utf8PathBuf;
@@ -36,15 +36,17 @@ pub async fn run_script(
     let contents = read_to_string(&path)
         .into_diagnostic()
         .wrap_err_with(|| format!("Failed to read script file {path}"))?;
-    let (mut options, mode) = parse_script(&contents, config)
+    let script = parse_script(&contents, config)
         .wrap_err_with(|| format!("Failed to parse script file {path}"))?;
-    let mode = mode_override.unwrap_or(mode);
+
+    let mode = mode_override.unwrap_or(script.mode);
+    let mut options = script.options;
 
     let target_path = Utf8PathBuf::from(format!("/{}", path.file_name().unwrap()));
-    let shell = if path.extension() == Some("nu") {
-        "nu"
-    } else {
-        "bash"
+    let shell = match (script.shell, path.extension()) {
+        (Some(shell), _) => shell,
+        (None, Some("nu")) => vec!["nu".into()],
+        _ => vec!["bash".into()],
     };
 
     options.mappings.push(
@@ -61,7 +63,10 @@ pub async fn run_script(
     cloud
         .exec(
             docker,
-            vec![shell.to_string(), target_path.into_string()],
+            shell
+                .into_iter()
+                .chain(once(target_path.into_string()))
+                .collect(),
             true,
             get_forward_env(),
         )
@@ -74,14 +79,22 @@ pub async fn run_script(
     Ok(())
 }
 
-fn parse_script(script: &str, config: &HazeConfig) -> Result<(CloudOptions, ScriptMode)> {
-    let (options, mode) = script
+#[derive(Debug)]
+struct Script {
+    options: CloudOptions,
+    mode: ScriptMode,
+    shell: Option<Vec<String>>,
+}
+
+fn parse_script(script: &str, config: &HazeConfig) -> Result<Script> {
+    let (options, mode, line_num) = script
         .lines()
-        .find_map(|line| line.strip_prefix("#! haze "))
+        .enumerate()
+        .find_map(|(i, line)| Some((i, line.strip_prefix("#! haze ")?)))
         .and_then(
-            |line| match (line.strip_prefix("shell"), line.strip_prefix("start")) {
-                (Some(shell_options), None) => Some((shell_options.trim(), ScriptMode::Shell)),
-                (None, Some(script_options)) => Some((script_options.trim(), ScriptMode::Start)),
+            |(i, line)| match (line.strip_prefix("shell"), line.strip_prefix("start")) {
+                (Some(shell_options), None) => Some((shell_options.trim(), ScriptMode::Shell, i)),
+                (None, Some(script_options)) => Some((script_options.trim(), ScriptMode::Start, i)),
                 _ => None,
             },
         )
@@ -90,5 +103,19 @@ fn parse_script(script: &str, config: &HazeConfig) -> Result<(CloudOptions, Scri
 
     let options =
         CloudOptions::parse(config, &mut options).wrap_err("Failed to parse shell options")?;
-    Ok((options, mode))
+
+    let first_content_line = script.lines().nth(line_num + 1).unwrap_or_default();
+    let shell = first_content_line.strip_prefix("#!").map(|bang| {
+        bang.split(" ")
+            .map(str::trim)
+            .filter(|s| !s.is_empty())
+            .map(String::from)
+            .collect()
+    });
+
+    Ok(Script {
+        options,
+        mode,
+        shell,
+    })
 }