Add env shorthand

crates/nu-command/src/default_context.rs

@@ -91,6 +91,7 @@ pub fn create_default_context() -> EngineState {
             Touch,
             Use,
             Where,
+            WithEnv,
             Wrap,
             Zip
         );

crates/nu-command/src/env/mod.rs

@@ -1,3 +1,5 @@
 mod let_env;
+mod with_env;
 
 pub use let_env::LetEnv;
+pub use with_env::WithEnv;

crates/nu-command/src/env/with_env.rs

@@ -0,0 +1,180 @@
+use std::{
+    collections::HashMap,
+    convert::{TryFrom, TryInto},
+};
+
+use nu_engine::{eval_block, CallExt};
+use nu_protocol::{
+    ast::Call,
+    engine::{Command, EngineState, Stack},
+    Example, PipelineData, ShellError, Signature, SyntaxShape, Value,
+};
+
+#[derive(Clone)]
+pub struct WithEnv;
+
+impl Command for WithEnv {
+    fn name(&self) -> &str {
+        "with-env"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("with-env")
+            .required(
+                "variable",
+                SyntaxShape::Any,
+                "the environment variable to temporarily set",
+            )
+            .required(
+                "block",
+                SyntaxShape::Block(Some(vec![SyntaxShape::Any])),
+                "the block to run once the variable is set",
+            )
+    }
+
+    fn usage(&self) -> &str {
+        "Runs a block with an environment variable set."
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        with_env(engine_state, stack, call, input)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Set the MYENV environment variable",
+                example: r#"with-env [MYENV "my env value"] { $nu.env.MYENV }"#,
+                result: Some(Value::test_string("my env value")),
+            },
+            Example {
+                description: "Set by primitive value list",
+                example: r#"with-env [X Y W Z] { $nu.env.X }"#,
+                result: Some(Value::test_string("Y")),
+            },
+            Example {
+                description: "Set by single row table",
+                example: r#"with-env [[X W]; [Y Z]] { $nu.env.W }"#,
+                result: Some(Value::test_string("Z")),
+            },
+            Example {
+                description: "Set by row(e.g. `open x.json` or `from json`)",
+                example: r#"echo '{"X":"Y","W":"Z"}'|from json|with-env $it { echo $nu.env.X $nu.env.W }"#,
+                result: None,
+            },
+        ]
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum EnvVar {
+    Proper(String),
+    Nothing,
+}
+
+impl TryFrom<&Value> for EnvVar {
+    type Error = ShellError;
+
+    fn try_from(value: &Value) -> Result<Self, Self::Error> {
+        if matches!(value, Value::Nothing { .. }) {
+            Ok(EnvVar::Nothing)
+        } else if let Ok(s) = value.as_string() {
+            if s.is_empty() {
+                Ok(EnvVar::Nothing)
+            } else {
+                Ok(EnvVar::Proper(s))
+            }
+        } else {
+            Err(ShellError::CantConvert("string".into(), value.span()?))
+        }
+    }
+}
+
+fn with_env(
+    engine_state: &EngineState,
+    stack: &mut Stack,
+    call: &Call,
+    input: PipelineData,
+) -> Result<PipelineData, ShellError> {
+    // let external_redirection = args.call_info.args.external_redirection;
+    let variable: Value = call.req(engine_state, stack, 0)?;
+
+    let block_id = call.positional[1]
+        .as_block()
+        .expect("internal error: expected block");
+    let block = engine_state.get_block(block_id).clone();
+    let mut stack = stack.collect_captures(&block.captures);
+
+    let mut env: HashMap<String, EnvVar> = HashMap::new();
+
+    match &variable {
+        Value::List { vals: table, .. } => {
+            if table.len() == 1 {
+                // single row([[X W]; [Y Z]])
+                match &table[0] {
+                    Value::Record { cols, vals, .. } => {
+                        for (k, v) in cols.iter().zip(vals.iter()) {
+                            env.insert(k.to_string(), v.try_into()?);
+                        }
+                    }
+                    _ => {
+                        return Err(ShellError::CantConvert(
+                            "string list or single row".into(),
+                            call.positional[1].span,
+                        ));
+                    }
+                }
+            } else {
+                // primitive values([X Y W Z])
+                for row in table.chunks(2) {
+                    if row.len() == 2 {
+                        env.insert(row[0].as_string()?, (&row[1]).try_into()?);
+                    }
+                }
+            }
+        }
+        // when get object by `open x.json` or `from json`
+        Value::Record { cols, vals, .. } => {
+            for (k, v) in cols.iter().zip(vals) {
+                env.insert(k.clone(), v.try_into()?);
+            }
+        }
+        _ => {
+            return Err(ShellError::CantConvert(
+                "string list or single row".into(),
+                call.positional[1].span,
+            ));
+        }
+    };
+
+    for (k, v) in env {
+        match v {
+            EnvVar::Nothing => {
+                stack.env_vars.remove(&k);
+            }
+            EnvVar::Proper(s) => {
+                stack.env_vars.insert(k, s);
+            }
+        }
+    }
+
+    eval_block(engine_state, &mut stack, &block, input)
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(WithEnv {})
+    }
+}

crates/nu-command/src/system/benchmark.rs

@@ -3,7 +3,7 @@ use std::time::Instant;
 use nu_engine::eval_block;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
-use nu_protocol::{PipelineData, Signature, SyntaxShape};
+use nu_protocol::{IntoPipelineData, PipelineData, Signature, SyntaxShape, Value};
 
 #[derive(Clone)]
 pub struct Benchmark;
@@ -42,7 +42,12 @@ impl Command for Benchmark {
         eval_block(engine_state, &mut stack, block, PipelineData::new())?.into_value();
 
         let end_time = Instant::now();
-        println!("{} ms", (end_time - start_time).as_millis());
-        Ok(PipelineData::new())
+
+        let output = Value::Duration {
+            val: (end_time - start_time).as_nanos() as i64,
+            span: call.head,
+        };
+
+        Ok(output.into_pipeline_data())
     }
 }

crates/nu-parser/src/parser.rs

@@ -627,27 +627,6 @@ pub fn parse_call(
     }
 
     let mut pos = 0;
-    let mut shorthand = vec![];
-
-    while pos < spans.len() {
-        // Check if there is any environment shorthand
-        let name = working_set.get_span_contents(spans[pos]);
-        let split: Vec<_> = name.splitn(2, |x| *x == b'=').collect();
-        if split.len() == 2 {
-            shorthand.push(split);
-            pos += 1;
-        } else {
-            break;
-        }
-    }
-
-    if pos == spans.len() {
-        return (
-            garbage(span(spans)),
-            Some(ParseError::UnknownCommand(spans[0])),
-        );
-    }
-
     let cmd_start = pos;
     let mut name_spans = vec![];
 
@@ -1681,6 +1660,54 @@ pub fn parse_string(
     }
 }
 
+pub fn parse_string_strict(
+    working_set: &mut StateWorkingSet,
+    span: Span,
+) -> (Expression, Option<ParseError>) {
+    let bytes = working_set.get_span_contents(span);
+    let (bytes, quoted) = if (bytes.starts_with(b"\"") && bytes.ends_with(b"\"") && bytes.len() > 1)
+        || (bytes.starts_with(b"\'") && bytes.ends_with(b"\'") && bytes.len() > 1)
+    {
+        (&bytes[1..(bytes.len() - 1)], true)
+    } else {
+        (bytes, false)
+    };
+
+    if let Ok(token) = String::from_utf8(bytes.into()) {
+        if quoted {
+            (
+                Expression {
+                    expr: Expr::String(token),
+                    span,
+                    ty: Type::String,
+                    custom_completion: None,
+                },
+                None,
+            )
+        } else if token.contains(' ') {
+            (
+                garbage(span),
+                Some(ParseError::Expected("string".into(), span)),
+            )
+        } else {
+            (
+                Expression {
+                    expr: Expr::String(token),
+                    span,
+                    ty: Type::String,
+                    custom_completion: None,
+                },
+                None,
+            )
+        }
+    } else {
+        (
+            garbage(span),
+            Some(ParseError::Expected("string".into(), span)),
+        )
+    }
+}
+
 //TODO: Handle error case for unknown shapes
 pub fn parse_shape_name(
     _working_set: &StateWorkingSet,
@@ -2963,12 +2990,126 @@ pub fn parse_expression(
     spans: &[Span],
     expand_aliases: bool,
 ) -> (Expression, Option<ParseError>) {
-    let bytes = working_set.get_span_contents(spans[0]);
+    let mut pos = 0;
+    let mut shorthand = vec![];
 
-    if is_math_expression_byte(bytes[0]) {
-        parse_math_expression(working_set, spans, None)
+    while pos < spans.len() {
+        // Check if there is any environment shorthand
+        let name = working_set.get_span_contents(spans[pos]);
+        let split = name.split(|x| *x == b'=');
+        let split: Vec<_> = split.collect();
+        if split.len() == 2 && !split[0].is_empty() {
+            let point = split[0].len() + 1;
+
+            let lhs = parse_string_strict(
+                working_set,
+                Span {
+                    start: spans[pos].start,
+                    end: spans[pos].start + point - 1,
+                },
+            );
+            let rhs = if spans[pos].start + point < spans[pos].end {
+                parse_string_strict(
+                    working_set,
+                    Span {
+                        start: spans[pos].start + point,
+                        end: spans[pos].end,
+                    },
+                )
+            } else {
+                (
+                    Expression {
+                        expr: Expr::String(String::new()),
+                        span: spans[pos],
+                        ty: Type::Nothing,
+                        custom_completion: None,
+                    },
+                    None,
+                )
+            };
+
+            if lhs.1.is_none() && rhs.1.is_none() {
+                shorthand.push((lhs.0, rhs.0));
+                pos += 1;
+            } else {
+                break;
+            }
+        } else {
+            break;
+        }
+    }
+
+    if pos == spans.len() {
+        return (
+            garbage(span(spans)),
+            Some(ParseError::UnknownCommand(spans[0])),
+        );
+    }
+
+    let bytes = working_set.get_span_contents(spans[pos]);
+
+    let (output, err) = if is_math_expression_byte(bytes[0]) {
+        parse_math_expression(working_set, &spans[pos..], None)
     } else {
-        parse_call(working_set, spans, expand_aliases)
+        parse_call(working_set, &spans[pos..], expand_aliases)
+    };
+
+    let with_env = working_set.find_decl(b"with-env");
+
+    if !shorthand.is_empty() {
+        if let Some(decl_id) = with_env {
+            let mut block = Block::default();
+            let ty = output.ty.clone();
+            block.stmts = vec![Statement::Pipeline(Pipeline {
+                expressions: vec![output],
+            })];
+
+            let mut seen = vec![];
+            let captures = find_captures_in_block(working_set, &block, &mut seen);
+            block.captures = captures;
+
+            let block_id = working_set.add_block(block);
+
+            let mut env_vars = vec![];
+            for sh in shorthand {
+                env_vars.push(sh.0);
+                env_vars.push(sh.1);
+            }
+
+            let positional = vec![
+                Expression {
+                    expr: Expr::List(env_vars),
+                    span: span(&spans[..pos]),
+                    ty: Type::Unknown,
+                    custom_completion: None,
+                },
+                Expression {
+                    expr: Expr::Block(block_id),
+                    span: span(&spans[pos..]),
+                    ty,
+                    custom_completion: None,
+                },
+            ];
+
+            (
+                Expression {
+                    expr: Expr::Call(Box::new(Call {
+                        head: span(spans),
+                        decl_id,
+                        named: vec![],
+                        positional,
+                    })),
+                    custom_completion: None,
+                    span: span(spans),
+                    ty: Type::Unknown,
+                },
+                err,
+            )
+        } else {
+            (output, err)
+        }
+    } else {
+        (output, err)
     }
 }
 

crates/nu-protocol/src/engine/stack.rs

@@ -64,6 +64,9 @@ impl Stack {
             }
         }
 
+        // FIXME: this is probably slow
+        output.env_vars = self.env_vars.clone();
+
         output
     }
 

crates/nu-protocol/src/value/mod.rs

@@ -199,8 +199,8 @@ impl Value {
             Value::Bool { val, .. } => val.to_string(),
             Value::Int { val, .. } => val.to_string(),
             Value::Float { val, .. } => val.to_string(),
-            Value::Filesize { val, .. } => format!("{} bytes", val),
-            Value::Duration { val, .. } => format!("{} ns", val),
+            Value::Filesize { val, .. } => format_filesize(val),
+            Value::Duration { val, .. } => format_duration(val),
             Value::Date { val, .. } => format!("{:?}", val),
             Value::Range { val, .. } => {
                 format!("{}..{}", val.from.into_string(), val.to.into_string())

src/tests.rs

@@ -804,10 +804,30 @@ fn help_works_with_missing_requirements() -> TestResult {
 
 #[test]
 fn scope_variable() -> TestResult {
-    run_test(r"let x = 3; $scope.vars.0", "$x")
+    run_test(r#"let x = 3; $scope.vars.0"#, "$x")
 }
 
 #[test]
 fn zip_ranges() -> TestResult {
-    run_test(r"1..3 | zip 4..6 | get 2.1", "6")
+    run_test(r#"1..3 | zip 4..6 | get 2.1"#, "6")
+}
+
+#[test]
+fn shorthand_env_1() -> TestResult {
+    run_test(r#"FOO=BAZ $nu.env.FOO"#, "BAZ")
+}
+
+#[test]
+fn shorthand_env_2() -> TestResult {
+    run_test(r#"FOO=BAZ FOO=MOO $nu.env.FOO"#, "MOO")
+}
+
+#[test]
+fn shorthand_env_3() -> TestResult {
+    run_test(r#"FOO=BAZ BAR=MOO $nu.env.FOO"#, "BAZ")
+}
+
+#[test]
+fn shorthand_env_4() -> TestResult {
+    fail_test(r#"FOO=BAZ FOO= $nu.env.FOO"#, "cannot find column")
 }