feat: Add the ability to set current dir and env vars

src/cmd.rs

@@ -7,22 +7,30 @@ use syn::{
 };
 
 pub struct Command {
+    cd: CurrentDir,
+    env_vars: EnvVars,
     program: Value,
     args: Option<Punctuated<LogicArg, Token![,]>>,
 }
 
 impl Parse for Command {
     fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let cd = input.parse()?;
+        let env_vars = input.parse()?;
         let program = input.parse()?;
 
         if input.is_empty() {
             Ok(Self {
+                cd,
+                env_vars,
                 program,
                 args: None,
             })
         } else {
             _ = input.parse::<Token![,]>()?;
             Ok(Self {
+                cd,
+                env_vars,
                 program,
                 args: Some(Punctuated::parse_terminated(input)?),
             })
@@ -32,22 +40,27 @@ impl Parse for Command {
 
 impl ToTokens for Command {
     fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
-        let Self { program, args } = self;
+        let Self {
+            cd,
+            env_vars,
+            program,
+            args,
+        } = self;
         let program = quote! { ::std::process::Command::new(#program) };
-        let args = args.as_ref().map(Punctuated::iter);
+        let args = args
+            .as_ref()
+            .map(Punctuated::iter)
+            .map_or_else(Vec::new, Iterator::collect);
 
-        tokens.extend(args.map_or_else(
-            || quote! { #program },
-            |args| {
-                quote! {
-                    {
-                        let mut _c = #program;
-                        #(#args)*
-                        _c
-                    }
-                }
-            },
-        ));
+        tokens.extend(quote! {
+            {
+                let mut _c = #program;
+                #cd
+                #env_vars
+                #(#args)*
+                _c
+            }
+        });
     }
 }
 
@@ -400,6 +413,99 @@ impl ToTokens for SingleArg {
     }
 }
 
+struct EnvVars(Option<Punctuated<EnvVar, Token![,]>>);
+
+impl Parse for EnvVars {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let fork = input.fork();
+        let ident = fork.cursor().ident();
+
+        match ident {
+            Some((ident, _)) if ident == "env" => {
+                _ = fork.parse::<syn::Ident>()?;
+                let envs;
+                braced!(envs in fork);
+                Punctuated::parse_terminated(&envs)
+                    .and_then(|envs| {
+                        _ = fork.parse::<Token![;]>()?;
+                        input.advance_to(&fork);
+                        Ok(Self(Some(envs)))
+                    })
+                    .or_else(|_| Ok(Self(None)))
+            }
+            _ => Ok(Self(None)),
+        }
+    }
+}
+
+impl ToTokens for EnvVars {
+    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+        let Self(envs) = self;
+        let envs = envs
+            .as_ref()
+            .map_or_else(Vec::new, |envs| envs.iter().collect());
+
+        tokens.extend(quote! {
+            #(#envs)*
+        });
+    }
+}
+
+struct EnvVar {
+    key: Value,
+    value: Value,
+}
+
+impl Parse for EnvVar {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let key = input.parse()?;
+        _ = input.parse::<Token![:]>()?;
+        let value = input.parse()?;
+
+        Ok(Self { key, value })
+    }
+}
+
+impl ToTokens for EnvVar {
+    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+        let Self { key, value } = self;
+
+        tokens.extend(quote! { _c.env(#key, #value); });
+    }
+}
+
+struct CurrentDir(Option<Value>);
+
+impl Parse for CurrentDir {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let fork = input.fork();
+        let ident = fork.cursor().ident();
+
+        match ident {
+            Some((ident, _)) if ident == "cd" => {
+                _ = fork.parse::<syn::Ident>();
+                fork.parse()
+                    .and_then(|value| {
+                        _ = fork.parse::<Token![;]>()?;
+                        input.advance_to(&fork);
+                        Ok(Self(Some(value)))
+                    })
+                    .or_else(|_| Ok(Self(None)))
+            }
+            _ => Ok(Self(None)),
+        }
+    }
+}
+
+impl ToTokens for CurrentDir {
+    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+        let Self(cd) = self;
+        let cd = cd.iter();
+
+        tokens.extend(quote! { #(_c.current_dir(#cd);)* });
+    }
+}
+
 enum Value {
     Lit(syn::Lit),
     Ident(syn::Ident),