add setenv and unsetenv builtins

2 files changed, 105 insertions, 32 deletions

diff --git a/src/env.rs b/src/env.rs
index e6cbe9c..221182e 100644
--- a/src/env.rs
+++ b/src/env.rs
@@ -96,6 +96,17 @@ impl Env {
         }
     }
 
+    pub fn set_vars(
+        &mut self,
+        it: impl Iterator<Item = (std::ffi::OsString, std::ffi::OsString)>,
+    ) {
+        match self {
+            Self::V0(env) => {
+                env.vars = it.collect();
+            }
+        }
+    }
+
     pub fn apply(&self, cmd: &mut pty_process::Command) {
         match self {
             Self::V0(env) => {
@@ -108,6 +119,7 @@ impl Env {
 
     pub fn update(&mut self) -> anyhow::Result<()> {
         self.set_current_dir(std::env::current_dir()?);
+        self.set_vars(std::env::vars_os());
         Ok(())
     }
 
diff --git a/src/pipeline/builtins/mod.rs b/src/pipeline/builtins/mod.rs
index e608ae7..e8d5e3d 100644
--- a/src/pipeline/builtins/mod.rs
+++ b/src/pipeline/builtins/mod.rs
@@ -17,6 +17,8 @@ static BUILTINS: once_cell::sync::Lazy<
 > = once_cell::sync::Lazy::new(|| {
     let mut builtins = std::collections::HashMap::new();
     builtins.insert("cd", &cd as Builtin);
+    builtins.insert("setenv", &setenv);
+    builtins.insert("unsetenv", &unsetenv);
     builtins.insert("echo", &echo);
     builtins.insert("and", &and);
     builtins.insert("or", &or);
@@ -25,6 +27,29 @@ static BUILTINS: once_cell::sync::Lazy<
     builtins
 });
 
+macro_rules! bail {
+    ($io:expr, $exe:expr, $msg:expr $(,)?) => {
+        $io.write_stderr(
+            format!("{}: {}\n", $exe.exe().display(), $msg).as_bytes()
+        )
+        .await
+        .unwrap();
+        return std::process::ExitStatus::from_raw(1 << 8);
+    };
+    ($io:expr, $exe:expr, $msg:expr, $($arg:tt)*) => {
+        $io.write_stderr(
+            format!("{}: ", $exe.exe().display()).as_bytes()
+        )
+        .await
+        .unwrap();
+        $io.write_stderr(format!($msg, $($arg)*).as_bytes())
+            .await
+            .unwrap();
+        $io.write_stderr(b"\n").await.unwrap();
+        return std::process::ExitStatus::from_raw(1 << 8);
+    };
+}
+
 // clippy can't tell that the type is necessary
 #[allow(clippy::unnecessary_wraps)]
 fn cd(
@@ -37,41 +62,12 @@ fn cd(
         _env: &Env,
         io: command::Io,
     ) -> std::process::ExitStatus {
-        macro_rules! bail {
-            ($msg:literal $(,)?) => {
-                io.write_stderr(format!("cd: {}\n", $msg).as_bytes())
-                    .await
-                    .unwrap();
-                return std::process::ExitStatus::from_raw(1 << 8);
-            };
-            ($msg:expr $(,)?) => {
-                io.write_stderr(format!("cd: {}\n", $msg).as_bytes())
-                    .await
-                    .unwrap();
-                return std::process::ExitStatus::from_raw(1 << 8);
-            };
-            ($msg:expr, $($arg:tt)*) => {
-                io.write_stderr(b"cd: ").await.unwrap();
-                io.write_stderr(format!($msg, $($arg)*).as_bytes())
-                    .await
-                    .unwrap();
-                io.write_stderr(b"\n").await.unwrap();
-                return std::process::ExitStatus::from_raw(1 << 8);
-            };
-        }
-
-        let dir = exe
-            .args()
-            .iter()
-            .map(std::convert::AsRef::as_ref)
-            .next()
-            .unwrap_or("");
-
+        let dir = exe.args().get(0).map_or("", String::as_str);
         let dir = if dir.is_empty() {
             if let Some(dir) = home(None) {
                 dir
             } else {
-                bail!("couldn't find current user");
+                bail!(io, exe, "couldn't find current user");
             }
         } else if dir.starts_with('~') {
             let path: std::path::PathBuf = dir.into();
@@ -88,6 +84,8 @@ fn cd(
                     home.join(path.strip_prefix(prefix).unwrap())
                 } else {
                     bail!(
+                        io,
+                        exe,
                         "no such user: {}",
                         name.map(std::ffi::OsStr::to_string_lossy)
                             .as_ref()
@@ -103,7 +101,13 @@ fn cd(
             dir.into()
         };
         if let Err(e) = std::env::set_current_dir(&dir) {
-            bail!("{}: {}", crate::format::io_error(&e), dir.display());
+            bail!(
+                io,
+                exe,
+                "{}: {}",
+                crate::format::io_error(&e),
+                dir.display()
+            );
         }
         async_std::process::ExitStatus::from_raw(0)
     }
@@ -113,6 +117,63 @@ fn cd(
     }))
 }
 
+#[allow(clippy::unnecessary_wraps)]
+fn setenv(
+    exe: crate::parse::Exe,
+    env: &Env,
+    io: command::Io,
+) -> anyhow::Result<command::Child> {
+    async fn async_setenv(
+        exe: crate::parse::Exe,
+        _env: &Env,
+        io: command::Io,
+    ) -> std::process::ExitStatus {
+        let k = if let Some(k) = exe.args().get(0).map(String::as_str) {
+            k
+        } else {
+            bail!(io, exe, "usage: setenv key value");
+        };
+        let v = if let Some(v) = exe.args().get(1).map(String::as_str) {
+            v
+        } else {
+            bail!(io, exe, "usage: setenv key value");
+        };
+
+        std::env::set_var(k, v);
+        async_std::process::ExitStatus::from_raw(0)
+    }
+
+    Ok(command::Child::new_fut(async move {
+        async_setenv(exe, env, io).await
+    }))
+}
+
+#[allow(clippy::unnecessary_wraps)]
+fn unsetenv(
+    exe: crate::parse::Exe,
+    env: &Env,
+    io: command::Io,
+) -> anyhow::Result<command::Child> {
+    async fn async_unsetenv(
+        exe: crate::parse::Exe,
+        _env: &Env,
+        io: command::Io,
+    ) -> std::process::ExitStatus {
+        let k = if let Some(k) = exe.args().get(0).map(String::as_str) {
+            k
+        } else {
+            bail!(io, exe, "usage: unsetenv key");
+        };
+
+        std::env::remove_var(k);
+        async_std::process::ExitStatus::from_raw(0)
+    }
+
+    Ok(command::Child::new_fut(async move {
+        async_unsetenv(exe, env, io).await
+    }))
+}
+
 // clippy can't tell that the type is necessary
 #[allow(clippy::unnecessary_wraps)]
 // mostly just for testing and ensuring that builtins work, i'll likely remove