diff --git a/crates/nu-command/src/filters/where_.rs b/crates/nu-command/src/filters/where_.rs index 691c04c345..28b55c7bf9 100644 --- a/crates/nu-command/src/filters/where_.rs +++ b/crates/nu-command/src/filters/where_.rs @@ -2,7 +2,8 @@ use nu_engine::{eval_block, CallExt}; use nu_protocol::ast::Call; use nu_protocol::engine::{CaptureBlock, Command, EngineState, Stack}; use nu_protocol::{ - Category, Example, IntoInterruptiblePipelineData, PipelineData, Signature, SyntaxShape, Value, + Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, ShellError, + Signature, Span, SyntaxShape, Value, }; #[derive(Clone)] @@ -19,7 +20,13 @@ impl Command for Where { fn signature(&self) -> nu_protocol::Signature { Signature::build("where") - .required("cond", SyntaxShape::RowCondition, "condition") + .optional("cond", SyntaxShape::RowCondition, "condition") + .named( + "block", + SyntaxShape::Block(Some(vec![SyntaxShape::Any])), + "use where with a block or variable instead", + Some('b'), + ) .category(Category::Filters) } @@ -34,51 +41,172 @@ impl Command for Where { call: &Call, input: PipelineData, ) -> Result { - let span = call.head; + if let Ok(Some(capture_block)) = call.get_flag::(engine_state, stack, "block") + { + let metadata = input.metadata(); + let ctrlc = engine_state.ctrlc.clone(); + let engine_state = engine_state.clone(); + let block = engine_state.get_block(capture_block.block_id).clone(); + let mut stack = stack.captures_to_stack(&capture_block.captures); + let orig_env_vars = stack.env_vars.clone(); + let orig_env_hidden = stack.env_hidden.clone(); + let span = call.head; + let redirect_stdout = call.redirect_stdout; + let redirect_stderr = call.redirect_stderr; - let metadata = input.metadata(); + match input { + PipelineData::Value(Value::Range { .. }, ..) + | PipelineData::Value(Value::List { .. }, ..) + | PipelineData::ListStream { .. } => Ok(input + .into_iter() + .filter_map(move |x| { + stack.with_env(&orig_env_vars, &orig_env_hidden); - let block: CaptureBlock = call.req(engine_state, stack, 0)?; - let mut stack = stack.captures_to_stack(&block.captures); - let block = engine_state.get_block(block.block_id).clone(); + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + stack.add_var(*var_id, x.clone()); + } + } - let ctrlc = engine_state.ctrlc.clone(); - let engine_state = engine_state.clone(); - - let redirect_stdout = call.redirect_stdout; - let redirect_stderr = call.redirect_stderr; - - Ok(input - .into_iter() - .filter_map(move |value| { - if let Some(var) = block.signature.get_positional(0) { - if let Some(var_id) = &var.var_id { - stack.add_var(*var_id, value.clone()); - } + match eval_block( + &engine_state, + &mut stack, + &block, + x.clone().into_pipeline_data(), + redirect_stdout, + redirect_stderr, + ) { + Ok(v) => { + if v.into_value(span).is_true() { + Some(x) + } else { + None + } + } + Err(error) => Some(Value::Error { error }), + } + }) + .into_pipeline_data(ctrlc)), + PipelineData::ExternalStream { stdout: None, .. } => { + Ok(PipelineData::new(call.head)) } - let result = eval_block( - &engine_state, - &mut stack, - &block, - PipelineData::new(span), - redirect_stdout, - redirect_stderr, - ); + PipelineData::ExternalStream { + stdout: Some(stream), + .. + } => Ok(stream + .into_iter() + .filter_map(move |x| { + stack.with_env(&orig_env_vars, &orig_env_hidden); - match result { - Ok(result) => { - let result = result.into_value(span); - if result.is_true() { - Some(value) - } else { - None + let x = match x { + Ok(x) => x, + Err(err) => return Some(Value::Error { error: err }), + }; + + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + stack.add_var(*var_id, x.clone()); + } + } + + match eval_block( + &engine_state, + &mut stack, + &block, + x.clone().into_pipeline_data(), + redirect_stdout, + redirect_stderr, + ) { + Ok(v) => { + if v.into_value(span).is_true() { + Some(x) + } else { + None + } + } + Err(error) => Some(Value::Error { error }), + } + }) + .into_pipeline_data(ctrlc)), + PipelineData::Value(x, ..) => { + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + stack.add_var(*var_id, x.clone()); } } - Err(err) => Some(Value::Error { error: err }), + Ok(match eval_block( + &engine_state, + &mut stack, + &block, + x.clone().into_pipeline_data(), + redirect_stdout, + redirect_stderr, + ) { + Ok(v) => { + if v.into_value(span).is_true() { + Some(x) + } else { + None + } + } + Err(error) => Some(Value::Error { error }), + } + .into_pipeline_data(ctrlc)) } - }) - .into_pipeline_data(ctrlc)) - .map(|x| x.set_metadata(metadata)) + } + .map(|x| x.set_metadata(metadata)) + } else { + let capture_block: Option = call.opt(engine_state, stack, 0)?; + if let Some(block) = capture_block { + let span = call.head; + + let metadata = input.metadata(); + let mut stack = stack.captures_to_stack(&block.captures); + let block = engine_state.get_block(block.block_id).clone(); + + let ctrlc = engine_state.ctrlc.clone(); + let engine_state = engine_state.clone(); + + let redirect_stdout = call.redirect_stdout; + let redirect_stderr = call.redirect_stderr; + Ok(input + .into_iter() + .filter_map(move |value| { + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + stack.add_var(*var_id, value.clone()); + } + } + let result = eval_block( + &engine_state, + &mut stack, + &block, + PipelineData::new(span), + redirect_stdout, + redirect_stderr, + ); + + match result { + Ok(result) => { + let result = result.into_value(span); + if result.is_true() { + Some(value) + } else { + None + } + } + Err(err) => Some(Value::Error { error: err }), + } + }) + .into_pipeline_data(ctrlc)) + .map(|x| x.set_metadata(metadata)) + } else { + Err(ShellError::MissingParameter( + "condition".to_string(), + call.head, + )) + } + } } fn examples(&self) -> Vec { @@ -103,6 +231,23 @@ impl Command for Where { example: "ls | where modified >= (date now) - 2wk", result: None, }, + Example { + description: "Get all numbers above 3 with an existing block condition", + example: "let a = {$in > 3}; [1, 2, 5, 6] | where -b $a", + result: Some(Value::List { + vals: vec![ + Value::Int { + val: 5, + span: Span::test_data(), + }, + Value::Int { + val: 6, + span: Span::test_data(), + }, + ], + span: Span::test_data(), + }), + }, ] } }