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, IntoPipelineData, PipelineData, Signature,
Span, SyntaxShape, Value,
};
#[derive(Clone)]
pub struct Each;
impl Command for Each {
fn name(&self) -> &str {
"each"
}
fn usage(&self) -> &str {
"Run a block on each element of input"
}
fn search_terms(&self) -> Vec<&str> {
vec!["for", "loop", "iterate"]
}
fn signature(&self) -> nu_protocol::Signature {
Signature::build("each")
.required(
"block",
SyntaxShape::Block(Some(vec![SyntaxShape::Any])),
"the block to run",
)
.switch("keep-empty", "keep empty result cells", Some('k'))
.switch("numbered", "iterate with an index", Some('n'))
.category(Category::Filters)
}
fn examples(&self) -> Vec<Example> {
let stream_test_1 = vec![
Value::Int {
val: 2,
span: Span::test_data(),
},
Value::Int {
val: 4,
span: Span::test_data(),
},
Value::Int {
val: 6,
span: Span::test_data(),
},
];
let stream_test_2 = vec![
Value::Nothing {
span: Span::test_data(),
},
Value::String {
val: "found 2!".to_string(),
span: Span::test_data(),
},
Value::Nothing {
span: Span::test_data(),
},
];
vec![
Example {
example: "[1 2 3] | each { |it| 2 * $it }",
description: "Multiplies elements in list",
result: Some(Value::List {
vals: stream_test_1,
span: Span::test_data(),
}),
},
Example {
example: r#"[1 2 3] | each { |it| if $it == 2 { echo "found 2!"} }"#,
description: "Iterate over each element, keeping only values that succeed",
result: Some(Value::List {
vals: vec![Value::String {
val: "found 2!".to_string(),
span: Span::test_data(),
}],
span: Span::test_data(),
}),
},
Example {
example: r#"[1 2 3] | each -n { |it| if $it.item == 2 { echo $"found 2 at ($it.index)!"} }"#,
description: "Iterate over each element, print the matching value and its index",
result: Some(Value::List {
vals: vec![Value::String {
val: "found 2 at 1!".to_string(),
span: Span::test_data(),
}],
span: Span::test_data(),
}),
},
Example {
example: r#"[1 2 3] | each --keep-empty { |it| if $it == 2 { echo "found 2!"} }"#,
description: "Iterate over each element, keeping all results",
result: Some(Value::List {
vals: stream_test_2,
span: Span::test_data(),
}),
},
]
}
fn run(
&self,
engine_state: &EngineState,
stack: &mut Stack,
call: &Call,
input: PipelineData,
) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
let capture_block: CaptureBlock = call.req(engine_state, stack, 0)?;
let numbered = call.has_flag("numbered");
let keep_empty = call.has_flag("keep-empty");
let metadata = input.metadata();
let ctrlc = engine_state.ctrlc.clone();
let outer_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;
match input {
PipelineData::Value(Value::Range { .. }, ..)
| PipelineData::Value(Value::List { .. }, ..)
| PipelineData::ListStream { .. } => Ok(input
.into_iter()
.enumerate()
.map(move |(idx, x)| {
stack.with_env(&orig_env_vars, &orig_env_hidden);
if let Some(var) = block.signature.get_positional(0) {
if let Some(var_id) = &var.var_id {
if numbered {
stack.add_var(
*var_id,
Value::Record {
cols: vec!["index".into(), "item".into()],
vals: vec![
Value::Int {
val: idx as i64,
span,
},
x.clone(),
],
span,
},
);
} else {
stack.add_var(*var_id, x.clone());
}
}
}
match eval_block(
&engine_state,
&mut stack,
&block,
x.into_pipeline_data(),
redirect_stdout,
redirect_stderr,
) {
Ok(v) => v.into_value(span),
Err(error) => Value::Error { error },
}
})
.into_pipeline_data(ctrlc)),
PipelineData::ExternalStream { stdout: None, .. } => Ok(PipelineData::new(call.head)),
PipelineData::ExternalStream {
stdout: Some(stream),
..
} => Ok(stream
.into_iter()
.enumerate()
.map(move |(idx, x)| {
stack.with_env(&orig_env_vars, &orig_env_hidden);
let x = match x {
Ok(x) => x,
Err(err) => return Value::Error { error: err },
};
if let Some(var) = block.signature.get_positional(0) {
if let Some(var_id) = &var.var_id {
if numbered {
stack.add_var(
*var_id,
Value::Record {
cols: vec!["index".into(), "item".into()],
vals: vec![
Value::Int {
val: idx as i64,
span,
},
x.clone(),
],
span,
},
);
} else {
stack.add_var(*var_id, x.clone());
}
}
}
match eval_block(
&engine_state,
&mut stack,
&block,
x.into_pipeline_data(),
redirect_stdout,
redirect_stderr,
) {
Ok(v) => v.into_value(span),
Err(error) => 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());
}
}
eval_block(
&engine_state,
&mut stack,
&block,
x.into_pipeline_data(),
redirect_stdout,
redirect_stderr,
)
}
}
.and_then(|x| {
x.filter(
move |x| if !keep_empty { !x.is_nothing() } else { true },
outer_ctrlc,
)
})
.map(|x| x.set_metadata(metadata))
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_examples() {
use crate::test_examples;
test_examples(Each {})
}
}