Add with-env command (#1717)

crates/nu-cli/src/cli.rs

@@ -261,6 +261,7 @@ pub fn create_default_context(
             whole_stream_command(Which),
             whole_stream_command(Debug),
             whole_stream_command(Alias),
+            whole_stream_command(WithEnv),
             // Statistics
             whole_stream_command(Size),
             whole_stream_command(Count),
@@ -856,8 +857,8 @@ async fn process_line(
             classified_block.block.expand_it_usage();
 
             trace!("{:#?}", classified_block);
-
+            let env = ctx.get_env();
-            match run_block(&classified_block.block, ctx, input_stream, &Scope::empty()).await {
+            match run_block(&classified_block.block, ctx, input_stream, &Scope::env(env)).await {
                 Ok(input) => {
                     // Running a pipeline gives us back a stream that we can then
                     // work through. At the top level, we just want to pull on the

crates/nu-cli/src/commands.rs

@@ -119,6 +119,7 @@ pub(crate) mod version;
 pub(crate) mod what;
 pub(crate) mod where_;
 pub(crate) mod which_;
+pub(crate) mod with_env;
 pub(crate) mod wrap;
 
 pub(crate) use autoview::Autoview;
@@ -241,4 +242,5 @@ pub(crate) use version::Version;
 pub(crate) use what::What;
 pub(crate) use where_::Where;
 pub(crate) use which_::Which;
+pub(crate) use with_env::WithEnv;
 pub(crate) use wrap::Wrap;

crates/nu-cli/src/commands/classified/external.rs

@@ -159,7 +159,7 @@ fn run_with_stdin(
         })
         .collect::<Vec<String>>();
 
-    spawn(&command, &path, &process_args[..], input, is_last)
+    spawn(&command, &path, &process_args[..], input, is_last, scope)
 }
 
 fn spawn(
@@ -168,6 +168,7 @@ fn spawn(
     args: &[String],
     input: InputStream,
     is_last: bool,
+    scope: &Scope,
 ) -> Result<InputStream, ShellError> {
     let command = command.clone();
 
@@ -197,6 +198,9 @@ fn spawn(
     process.current_dir(path);
     trace!(target: "nu::run::external", "cwd = {:?}", &path);
 
+    process.env_clear();
+    process.envs(scope.env.iter());
+
     // We want stdout regardless of what
     // we are doing ($it case or pipe stdin)
     if !is_last {

crates/nu-cli/src/commands/classified/internal.rs

@@ -33,6 +33,7 @@ pub(crate) fn run_internal_command(
 
     let mut result = trace_out_stream!(target: "nu::trace_stream::internal", "output" = result);
     let mut context = context.clone();
+    let scope = scope.clone();
 
     let stream = async_stream! {
         let mut soft_errs: Vec<ShellError> = vec![];
@@ -67,7 +68,7 @@ pub(crate) fn run_internal_command(
                                         is_last: false,
                                     },
                                     name_tag: Tag::unknown_anchor(command.name_span),
-                                    scope: Scope::empty(),
+                                    scope: scope.clone(),
                                 }
                             };
                             let mut result = converter.run(new_args.with_input(vec![tagged_contents]), &context.registry);

crates/nu-cli/src/commands/merge.rs

@@ -6,7 +6,7 @@ use crate::prelude::*;
 
 use indexmap::IndexMap;
 use nu_errors::ShellError;
-use nu_protocol::{hir::Block, ReturnSuccess, Scope, Signature, SyntaxShape, UntaggedValue, Value};
+use nu_protocol::{hir::Block, ReturnSuccess, Signature, SyntaxShape, UntaggedValue, Value};
 pub struct Merge;
 
 #[derive(Deserialize)]
@@ -48,6 +48,7 @@ fn merge(
     let block = merge_args.block;
     let registry = context.registry.clone();
     let mut input = context.input;
+    let scope = raw_args.call_info.scope.clone();
 
     let mut context = Context::from_raw(&raw_args, &registry);
 
@@ -55,7 +56,7 @@ fn merge(
         let table: Option<Vec<Value>> = match run_block(&block,
                 &mut context,
                 InputStream::empty(),
-                &Scope::empty()).await {
+                &scope).await {
             Ok(mut stream) => Some(stream.drain_vec().await),
             Err(err) => {
                 yield Err(err);

crates/nu-cli/src/commands/save.rs

@@ -1,7 +1,7 @@
 use crate::commands::{UnevaluatedCallInfo, WholeStreamCommand};
 use crate::prelude::*;
 use nu_errors::ShellError;
-use nu_protocol::{Primitive, ReturnSuccess, Scope, Signature, SyntaxShape, UntaggedValue, Value};
+use nu_protocol::{Primitive, ReturnSuccess, Signature, SyntaxShape, UntaggedValue, Value};
 use nu_source::Tagged;
 use std::path::{Path, PathBuf};
 
@@ -175,6 +175,7 @@ fn save(
 ) -> Result<OutputStream, ShellError> {
     let mut full_path = PathBuf::from(shell_manager.path());
     let name_tag = name.clone();
+    let scope = raw_args.call_info.scope.clone();
 
     let stream = async_stream! {
         let input: Vec<Value> = input.collect().await;
@@ -236,7 +237,7 @@ fn save(
                                     is_last: false,
                                 },
                                 name_tag: raw_args.call_info.name_tag,
-                                scope: Scope::empty(), // FIXME?
+                                scope,
                             }
                         };
                         let mut result = converter.run(new_args.with_input(input), &registry);

crates/nu-cli/src/commands/with_env.rs

@@ -0,0 +1,87 @@
+use crate::commands::classified::block::run_block;
+use crate::commands::WholeStreamCommand;
+use crate::prelude::*;
+use nu_errors::ShellError;
+use nu_protocol::{hir::Block, ReturnSuccess, Signature, SyntaxShape};
+use nu_source::Tagged;
+
+pub struct WithEnv;
+
+#[derive(Deserialize, Debug)]
+struct WithEnvArgs {
+    variable: (Tagged<String>, Tagged<String>),
+    block: Block,
+}
+impl WholeStreamCommand 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,
+                "the block to run once the variable is set",
+            )
+    }
+
+    fn usage(&self) -> &str {
+        "Runs a block with an environment set. Eg) with-env [NAME 'foo'] { echo $nu.env.NAME }"
+    }
+
+    fn run(
+        &self,
+        args: CommandArgs,
+        registry: &CommandRegistry,
+    ) -> Result<OutputStream, ShellError> {
+        Ok(args.process_raw(registry, with_env)?.run())
+    }
+}
+
+fn with_env(
+    WithEnvArgs { variable, block }: WithEnvArgs,
+    context: RunnableContext,
+    raw_args: RawCommandArgs,
+) -> Result<OutputStream, ShellError> {
+    let scope = raw_args
+        .call_info
+        .scope
+        .clone()
+        .set_env_var(variable.0.item, variable.1.item);
+    let registry = context.registry.clone();
+    let input = context.input;
+    let mut context = Context::from_raw(&raw_args, &registry);
+
+    let stream = async_stream! {
+        let result = run_block(
+            &block,
+            &mut context,
+            input,
+            &scope.clone(),
+        ).await;
+
+        match result {
+            Ok(mut stream) => {
+                while let Some(result) = stream.next().await {
+                    yield Ok(ReturnSuccess::Value(result));
+                }
+
+                let errors = context.get_errors();
+                if let Some(error) = errors.first() {
+                    yield Err(error.clone());
+                }
+            }
+            Err(e) => {
+                yield Err(e);
+            }
+        }
+    };
+
+    Ok(stream.to_output_stream())
+}

crates/nu-cli/src/context.rs

@@ -258,4 +258,12 @@ impl Context {
             input,
         }
     }
+
+    pub fn get_env(&self) -> IndexMap<String, String> {
+        let mut output = IndexMap::new();
+        for (var, value) in self.host.lock().vars() {
+            output.insert(var, value);
+        }
+        output
+    }
 }

crates/nu-cli/src/evaluate/evaluator.rs

@@ -151,7 +151,7 @@ fn evaluate_reference(name: &hir::Variable, scope: &Scope, tag: Tag) -> Result<V
     match name {
         hir::Variable::It(_) => Ok(scope.it.value.clone().into_value(tag)),
         hir::Variable::Other(name, _) => match name {
-            x if x == "$nu" => crate::evaluate::variables::nu(tag),
+            x if x == "$nu" => crate::evaluate::variables::nu(scope, tag),
             x if x == "$true" => Ok(Value {
                 value: UntaggedValue::boolean(true),
                 tag,

crates/nu-cli/src/evaluate/variables.rs

@@ -1,15 +1,15 @@
 use crate::cli::History;
 use nu_errors::ShellError;
-use nu_protocol::{TaggedDictBuilder, UntaggedValue, Value};
+use nu_protocol::{Scope, TaggedDictBuilder, UntaggedValue, Value};
 use nu_source::Tag;
 
-pub fn nu(tag: impl Into<Tag>) -> Result<Value, ShellError> {
+pub fn nu(scope: &Scope, tag: impl Into<Tag>) -> Result<Value, ShellError> {
     let tag = tag.into();
 
     let mut nu_dict = TaggedDictBuilder::new(&tag);
 
     let mut dict = TaggedDictBuilder::new(&tag);
-    for v in std::env::vars() {
+    for v in scope.env.iter() {
         if v.0 != "PATH" && v.0 != "Path" {
             dict.insert_untagged(v.0, UntaggedValue::string(v.1));
         }

crates/nu-cli/tests/commands/mod.rs

@@ -46,4 +46,5 @@ mod touch;
 mod trim;
 mod uniq;
 mod where_;
+mod with_env;
 mod wrap;

crates/nu-cli/tests/commands/with_env.rs

@@ -0,0 +1,11 @@
+use nu_test_support::nu;
+
+#[test]
+fn with_env_extends_environment() {
+    let actual = nu!(
+        cwd: "tests/fixtures/formats",
+        "with-env [FOO BARRRR] {echo $nu.env} | get FOO"
+    );
+
+    assert_eq!(actual, "BARRRR");
+}

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

@@ -10,6 +10,7 @@ use std::fmt::Debug;
 pub struct Scope {
     pub it: Value,
     pub vars: IndexMap<String, Value>,
+    pub env: IndexMap<String, String>,
 }
 
 impl Scope {
@@ -18,6 +19,7 @@ impl Scope {
         Scope {
             it,
             vars: IndexMap::new(),
+            env: IndexMap::new(),
         }
     }
 }
@@ -28,6 +30,7 @@ impl Scope {
         Scope {
             it: UntaggedValue::Primitive(Primitive::Nothing).into_untagged_value(),
             vars: IndexMap::new(),
+            env: IndexMap::new(),
         }
     }
 
@@ -36,6 +39,15 @@ impl Scope {
         Scope {
             it: value,
             vars: IndexMap::new(),
+            env: IndexMap::new(),
+        }
+    }
+
+    pub fn env(env: IndexMap<String, String>) -> Scope {
+        Scope {
+            it: UntaggedValue::Primitive(Primitive::Nothing).into_untagged_value(),
+            vars: IndexMap::new(),
+            env,
         }
     }
 
@@ -43,6 +55,7 @@ impl Scope {
         Scope {
             it: value,
             vars: self.vars,
+            env: self.env,
         }
     }
 
@@ -52,6 +65,17 @@ impl Scope {
         Scope {
             it: self.it,
             vars: new_vars,
+            env: self.env,
+        }
+    }
+
+    pub fn set_env_var(self, variable: String, value: String) -> Scope {
+        let mut new_env_vars = self.env.clone();
+        new_env_vars.insert(variable, value);
+        Scope {
+            it: self.it,
+            vars: self.vars,
+            env: new_env_vars,
         }
     }
 }