From 1949ba080e3f1f33e04efaff62f33f6eb2bc8b9d Mon Sep 17 00:00:00 2001 From: JT Date: Thu, 4 Nov 2021 15:32:35 +1300 Subject: [PATCH] Add env shorthand --- crates/nu-command/src/default_context.rs | 1 + crates/nu-command/src/env/mod.rs | 2 + crates/nu-command/src/env/with_env.rs | 180 ++++++++++++++++++++ crates/nu-command/src/system/benchmark.rs | 11 +- crates/nu-parser/src/parser.rs | 191 +++++++++++++++++++--- crates/nu-protocol/src/engine/stack.rs | 3 + crates/nu-protocol/src/value/mod.rs | 4 +- src/tests.rs | 24 ++- 8 files changed, 384 insertions(+), 32 deletions(-) create mode 100644 crates/nu-command/src/env/with_env.rs diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 95635f093c..116b33b3cf 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -91,6 +91,7 @@ pub fn create_default_context() -> EngineState { Touch, Use, Where, + WithEnv, Wrap, Zip ); diff --git a/crates/nu-command/src/env/mod.rs b/crates/nu-command/src/env/mod.rs index fa42c1b093..48e8293dfd 100644 --- a/crates/nu-command/src/env/mod.rs +++ b/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; diff --git a/crates/nu-command/src/env/with_env.rs b/crates/nu-command/src/env/with_env.rs new file mode 100644 index 0000000000..ae3d3a3776 --- /dev/null +++ b/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 { + with_env(engine_state, stack, call, input) + } + + fn examples(&self) -> Vec { + 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 { + 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 { + // 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 = 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 {}) + } +} diff --git a/crates/nu-command/src/system/benchmark.rs b/crates/nu-command/src/system/benchmark.rs index 7b99c2009a..e786c40885 100644 --- a/crates/nu-command/src/system/benchmark.rs +++ b/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()) } } diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs index 25be7f6d79..3d932b4c89 100644 --- a/crates/nu-parser/src/parser.rs +++ b/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) { + 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) { - 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) } } diff --git a/crates/nu-protocol/src/engine/stack.rs b/crates/nu-protocol/src/engine/stack.rs index 45b013de59..f2fc9fc0ac 100644 --- a/crates/nu-protocol/src/engine/stack.rs +++ b/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 } diff --git a/crates/nu-protocol/src/value/mod.rs b/crates/nu-protocol/src/value/mod.rs index b6ea1e2768..84ff1dd51d 100644 --- a/crates/nu-protocol/src/value/mod.rs +++ b/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()) diff --git a/src/tests.rs b/src/tests.rs index 41595e65c2..a9042eb06d 100644 --- a/src/tests.rs +++ b/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") }