From ffad7953bba4da39dc17ec06fb97e35eff62a6a1 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Fri, 5 Apr 2024 11:41:51 +0800
Subject: [PATCH] fix: load dotenv before handle @env (#308)

---
 src/build.rs       |  8 ++++----
 src/command/mod.rs | 13 +++++++------
 src/matcher.rs     | 13 +++++++++----
 src/utils.rs       | 20 ++++++++++++++++++--
 4 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/src/build.rs b/src/build.rs
index da7b94c..3386d35 100644
--- a/src/build.rs
+++ b/src/build.rs
@@ -1,7 +1,7 @@
 use crate::{
     command::Command,
     param::{FlagOptionParam, Param},
-    utils::{escape_shell_words, expand_dotenv, META_DOTENV},
+    utils::{escape_shell_words, expand_dotenv},
     ChoiceValue, DefaultValue,
 };
 use anyhow::Result;
@@ -193,7 +193,7 @@ pub fn build(source: &str, root_name: &str) -> Result<String> {
 
 fn build_root(cmd: &Command) -> String {
     let command = build_command(cmd);
-    let dotenv = if let Some(value) = cmd.get_metadata(META_DOTENV) {
+    let dotenv = if let Some(value) = cmd.dotenv() {
         format!("\n    {}", expand_dotenv(value))
     } else {
         String::new()
@@ -224,8 +224,8 @@ _argc_run() {{
     argc__args=("$(basename "$0" .sh)" "$@")
     argc__positionals=()
     _argc_index=1
-    _argc_len="${{#argc__args[@]}}"
-    _argc_parse{dotenv}{before_hook}
+    _argc_len="${{#argc__args[@]}}"{dotenv}
+    _argc_parse{before_hook}
     if [ -n "$argc__fn" ]; then
         $argc__fn "${{argc__positionals[@]}}"{after_hook}
     fi
diff --git a/src/command/mod.rs b/src/command/mod.rs
index f2ea8d3..aad54ce 100644
--- a/src/command/mod.rs
+++ b/src/command/mod.rs
@@ -100,12 +100,7 @@ impl Command {
             if self.get_metadata(META_COMBINE_SHORTS).is_some() {
                 extra.insert("combine_shorts".into(), true.into());
             }
-            if let Some(dotenv) = self.get_metadata(META_DOTENV) {
-                let dotenv = if dotenv.is_empty() {
-                    ".env".to_string()
-                } else {
-                    dotenv.to_string()
-                };
+            if let Some(dotenv) = self.dotenv() {
                 extra.insert("dotenv".into(), dotenv.into());
             }
             let (before_hook, after_hook) = self.exist_hooks();
@@ -611,6 +606,12 @@ impl Command {
             && self.positional_params[0].terminated()
     }
 
+    pub(crate) fn dotenv(&self) -> Option<&str> {
+        let dotenv = self.get_metadata(META_DOTENV)?;
+        let dotenv = if dotenv.is_empty() { ".env" } else { dotenv };
+        Some(dotenv)
+    }
+
     fn update_recursively(&mut self, paths: Vec<String>) {
         self.paths = paths.clone();
 
diff --git a/src/matcher.rs b/src/matcher.rs
index 4eb4228..6e2547d 100644
--- a/src/matcher.rs
+++ b/src/matcher.rs
@@ -10,7 +10,7 @@ use crate::{
     command::{Command, SymbolParam},
     compgen::CompColor,
     param::{ChoiceValue, FlagOptionParam, Param, ParamData, PositionalParam},
-    utils::{argc_var_name, run_param_fns, META_COMBINE_SHORTS, META_DOTENV},
+    utils::{argc_var_name, load_dotenv, run_param_fns, META_COMBINE_SHORTS},
     Shell,
 };
 
@@ -292,10 +292,15 @@ impl<'a, 'b> Matcher<'a, 'b> {
             }
         }
 
-        let mut envs = HashMap::new();
         let last_cmd = *cmds.last().unwrap();
+
+        let mut envs = HashMap::new();
+        let dotenv_vars = root_cmd.dotenv().and_then(load_dotenv).unwrap_or_default();
         for param in &last_cmd.env_params {
-            if let Ok(value) = std::env::var(param.id()) {
+            if let Some(value) = std::env::var(param.id())
+                .ok()
+                .or_else(|| dotenv_vars.get(param.id()).cloned())
+            {
                 envs.insert(param.id(), value);
             }
             add_param_choice_fn(&mut choice_fns, param)
@@ -476,7 +481,7 @@ impl<'a, 'b> Matcher<'a, 'b> {
         let level = cmds_len - 1;
         let last_cmd = self.cmds[level];
 
-        if let Some(value) = root_cmd.get_metadata(META_DOTENV) {
+        if let Some(value) = root_cmd.dotenv() {
             output.push(ArgcValue::Dotenv(value.to_string()))
         }
 
diff --git a/src/utils.rs b/src/utils.rs
index 5980596..627e12e 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,7 +1,7 @@
 use convert_case::{Boundary, Converter, Pattern};
 use std::{
     collections::HashMap,
-    env,
+    env, fs,
     path::{Path, PathBuf},
     process, thread,
 };
@@ -186,7 +186,7 @@ pub fn path_env_with_exe() -> String {
 }
 
 pub fn expand_dotenv(value: &str) -> String {
-    let value = if value.is_empty() { ".env" } else { value };
+    let value = escape_shell_words(value);
     format!("[ -f {value} ] && set -o allexport && . {value} && set +o allexport")
 }
 
@@ -198,6 +198,22 @@ pub fn argc_var_name(id: &str) -> String {
     format!("{VARIABLE_PREFIX}{}", sanitize_var_name(id))
 }
 
+pub fn load_dotenv(path: &str) -> Option<HashMap<String, String>> {
+    let contents = fs::read_to_string(path).ok()?;
+    let mut output = HashMap::new();
+    for line in contents.lines() {
+        if line.starts_with('#') || line.trim().is_empty() {
+            continue;
+        }
+        if let Some((key, value)) = line.split_once('=') {
+            let key = key.trim().to_string();
+            let value = value.trim().to_string();
+            output.insert(key, value);
+        }
+    }
+    Some(output)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;