Add env() function (#1613)

src/assignment_resolver.rs

@@ -73,6 +73,13 @@ impl<'src: 'run, 'run> AssignmentResolver<'src, 'run> {
       Expression::Call { thunk } => match thunk {
         Thunk::Nullary { .. } => Ok(()),
         Thunk::Unary { arg, .. } => self.resolve_expression(arg),
+        Thunk::UnaryOpt { args: (a, b), .. } => {
+          self.resolve_expression(a)?;
+          if let Some(b) = b.as_ref() {
+            self.resolve_expression(b)?;
+          }
+          Ok(())
+        }
         Thunk::Binary { args: [a, b], .. } => {
           self.resolve_expression(a)?;
           self.resolve_expression(b)

src/evaluator.rs

@@ -92,6 +92,23 @@ impl<'src, 'run> Evaluator<'src, 'run> {
               message,
             }
           }),
+          UnaryOpt {
+            name,
+            function,
+            args: (a, b),
+            ..
+          } => {
+            let a = self.evaluate_expression(a)?;
+            let b = match b.as_ref() {
+              Some(b) => Some(self.evaluate_expression(b)?),
+              None => None,
+            };
+
+            function(&context, &a, b.as_deref()).map_err(|message| Error::FunctionCall {
+              function: *name,
+              message,
+            })
+          }
           Binary {
             name,
             function,

src/function.rs

@@ -10,6 +10,7 @@ use {
 pub(crate) enum Function {
   Nullary(fn(&FunctionContext) -> Result<String, String>),
   Unary(fn(&FunctionContext, &str) -> Result<String, String>),
+  UnaryOpt(fn(&FunctionContext, &str, Option<&str>) -> Result<String, String>),
   Binary(fn(&FunctionContext, &str, &str) -> Result<String, String>),
   BinaryPlus(fn(&FunctionContext, &str, &str, &[String]) -> Result<String, String>),
   Ternary(fn(&FunctionContext, &str, &str, &str) -> Result<String, String>),
@@ -21,6 +22,7 @@ pub(crate) fn get(name: &str) -> Option<Function> {
     "arch" => Nullary(arch),
     "capitalize" => Unary(capitalize),
     "clean" => Unary(clean),
+    "env" => UnaryOpt(env),
     "env_var" => Unary(env_var),
     "env_var_or_default" => Binary(env_var_or_default),
     "error" => Unary(error),
@@ -70,6 +72,7 @@ impl Function {
     match *self {
       Nullary(_) => 0..0,
       Unary(_) => 1..1,
+      UnaryOpt(_) => 1..2,
       Binary(_) => 2..2,
       BinaryPlus(_) => 2..usize::MAX,
       Ternary(_) => 3..3,
@@ -144,6 +147,13 @@ fn env_var_or_default(
   }
 }
 
+fn env(context: &FunctionContext, key: &str, default: Option<&str>) -> Result<String, String> {
+  match default {
+    Some(val) => env_var_or_default(context, key, val),
+    None => env_var(context, key),
+  }
+}
+
 fn error(_context: &FunctionContext, message: &str) -> Result<String, String> {
   Err(message.to_owned())
 }

src/node.rs

@@ -79,6 +79,15 @@ impl<'src> Node<'src> for Expression<'src> {
             tree.push_mut(name.lexeme());
             tree.push_mut(arg.tree());
           }
+          UnaryOpt {
+            name, args: (a, b), ..
+          } => {
+            tree.push_mut(name.lexeme());
+            tree.push_mut(a.tree());
+            if let Some(b) = b.as_ref() {
+              tree.push_mut(b.tree());
+            }
+          }
           Binary {
             name, args: [a, b], ..
           } => {

src/parser.rs

@@ -2333,6 +2333,34 @@ mod tests {
     },
   }
 
+  error! {
+    name: function_argument_count_too_high_unary_opt,
+    input: "x := env('foo', 'foo', 'foo')",
+    offset: 5,
+    line: 0,
+    column: 5,
+    width: 3,
+    kind: FunctionArgumentCountMismatch {
+      function: "env",
+      found: 3,
+      expected: 1..2,
+    },
+  }
+
+  error! {
+    name: function_argument_count_too_low_unary_opt,
+    input: "x := env()",
+    offset: 5,
+    line: 0,
+    column: 5,
+    width: 3,
+    kind: FunctionArgumentCountMismatch {
+      function: "env",
+      found: 0,
+      expected: 1..2,
+    },
+  }
+
   error! {
     name: function_argument_count_binary,
     input: "x := env_var_or_default('foo')",

src/summary.rs

@@ -225,6 +225,23 @@ impl Expression {
           name: name.lexeme().to_owned(),
           arguments: vec![Expression::new(arg)],
         },
+        full::Thunk::UnaryOpt {
+          name,
+          args: (a, opt_b),
+          ..
+        } => {
+          let mut arguments = vec![];
+
+          if let Some(b) = opt_b.as_ref() {
+            arguments.push(Expression::new(b));
+          }
+
+          arguments.push(Expression::new(a));
+          Expression::Call {
+            name: name.lexeme().to_owned(),
+            arguments,
+          }
+        }
         full::Thunk::Binary {
           name, args: [a, b], ..
         } => Expression::Call {

src/thunk.rs

@@ -14,6 +14,12 @@ pub(crate) enum Thunk<'src> {
     function: fn(&FunctionContext, &str) -> Result<String, String>,
     arg: Box<Expression<'src>>,
   },
+  UnaryOpt {
+    name: Name<'src>,
+    #[derivative(Debug = "ignore", PartialEq = "ignore")]
+    function: fn(&FunctionContext, &str, Option<&str>) -> Result<String, String>,
+    args: (Box<Expression<'src>>, Box<Option<Expression<'src>>>),
+  },
   Binary {
     name: Name<'src>,
     #[derivative(Debug = "ignore", PartialEq = "ignore")]
@@ -39,6 +45,7 @@ impl<'src> Thunk<'src> {
     match self {
       Self::Nullary { name, .. }
       | Self::Unary { name, .. }
+      | Self::UnaryOpt { name, .. }
       | Self::Binary { name, .. }
       | Self::BinaryPlus { name, .. }
       | Self::Ternary { name, .. } => name,
@@ -60,6 +67,18 @@ impl<'src> Thunk<'src> {
           arg: Box::new(arguments.pop().unwrap()),
           name,
         }),
+        (Function::UnaryOpt(function), 1..=2) => {
+          let a = Box::new(arguments.remove(0));
+          let b = match arguments.pop() {
+            Some(value) => Box::new(Some(value)),
+            None => Box::new(None),
+          };
+          Ok(Thunk::UnaryOpt {
+            function,
+            args: (a, b),
+            name,
+          })
+        }
         (Function::Binary(function), 2) => {
           let b = Box::new(arguments.pop().unwrap());
           let a = Box::new(arguments.pop().unwrap());
@@ -105,6 +124,15 @@ impl Display for Thunk<'_> {
     match self {
       Nullary { name, .. } => write!(f, "{}()", name.lexeme()),
       Unary { name, arg, .. } => write!(f, "{}({arg})", name.lexeme()),
+      UnaryOpt {
+        name, args: (a, b), ..
+      } => {
+        if let Some(b) = b.as_ref() {
+          write!(f, "{}({a}, {b})", name.lexeme())
+        } else {
+          write!(f, "{}({a})", name.lexeme())
+        }
+      }
       Binary {
         name, args: [a, b], ..
       } => write!(f, "{}({a}, {b})", name.lexeme()),
@@ -139,6 +167,14 @@ impl<'src> Serialize for Thunk<'src> {
     match self {
       Self::Nullary { .. } => {}
       Self::Unary { arg, .. } => seq.serialize_element(&arg)?,
+      Self::UnaryOpt {
+        args: (a, opt_b), ..
+      } => {
+        seq.serialize_element(a)?;
+        if let Some(b) = opt_b.as_ref() {
+          seq.serialize_element(b)?;
+        }
+      }
       Self::Binary { args, .. } => {
         for arg in args {
           seq.serialize_element(arg)?;

src/variables.rs

@@ -20,6 +20,14 @@ impl<'expression, 'src> Iterator for Variables<'expression, 'src> {
         Expression::Call { thunk } => match thunk {
           Thunk::Nullary { .. } => {}
           Thunk::Unary { arg, .. } => self.stack.push(arg),
+          Thunk::UnaryOpt {
+            args: (a, opt_b), ..
+          } => {
+            self.stack.push(a);
+            if let Some(b) = opt_b.as_ref() {
+              self.stack.push(b);
+            }
+          }
           Thunk::Binary { args, .. } => {
             for arg in args.iter().rev() {
               self.stack.push(arg);

tests/misc.rs

@@ -1914,6 +1914,66 @@ test! {
   shell: false,
 }
 
+test! {
+  name: env_function_as_env_var,
+  justfile: "
+    foo: (bar env('x'))
+
+    bar arg:
+      echo {{arg}}
+  ",
+  args: (),
+  env: { "x": "z", },
+  stdout: "z\n",
+  stderr: "echo z\n",
+  shell: false,
+}
+
+test! {
+  name: env_function_as_env_var_or_default,
+  justfile: "
+    foo: (bar env('x', 'y'))
+
+    bar arg:
+      echo {{arg}}
+  ",
+  args: (),
+  env: { "x": "z", },
+  stdout: "z\n",
+  stderr: "echo z\n",
+  shell: false,
+}
+
+test! {
+  name: env_function_as_env_var_with_existing_env_var,
+  justfile: "
+    foo: (bar env('x'))
+
+    bar arg:
+      echo {{arg}}
+  ",
+  args: (),
+  env: { "x": "z", },
+  stdout: "z\n",
+  stderr: "echo z\n",
+  shell: false,
+}
+
+test! {
+  name: env_function_as_env_var_or_default_with_existing_env_var,
+  justfile: "
+    foo: (bar env('x', 'y'))
+
+    bar arg:
+      echo {{arg}}
+  ",
+  args: (),
+  env: { "x": "z", },
+  stdout: "z\n",
+  stderr: "echo z\n",
+  shell: false,
+}
+
 test! {
   name: dependency_argument_backtick,
   justfile: "