diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 83bf2e3951..18ac7e1748 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -97,6 +97,7 @@ pub fn create_default_context(cwd: impl AsRef) -> EngineState { Transpose, Uniq, Update, + UpdateCells, Where, Wrap, Zip, diff --git a/crates/nu-command/src/filters/mod.rs b/crates/nu-command/src/filters/mod.rs index 9da0d7ab50..31dd2fa1c9 100644 --- a/crates/nu-command/src/filters/mod.rs +++ b/crates/nu-command/src/filters/mod.rs @@ -36,6 +36,7 @@ mod sort_by; mod transpose; mod uniq; mod update; +mod update_cells; mod where_; mod wrap; mod zip_; @@ -78,6 +79,7 @@ pub use sort_by::SortBy; pub use transpose::Transpose; pub use uniq::*; pub use update::Update; +pub use update_cells::UpdateCells; pub use where_::Where; pub use wrap::Wrap; pub use zip_::Zip; diff --git a/crates/nu-command/src/filters/update_cells.rs b/crates/nu-command/src/filters/update_cells.rs new file mode 100644 index 0000000000..6dbe50c0c6 --- /dev/null +++ b/crates/nu-command/src/filters/update_cells.rs @@ -0,0 +1,247 @@ +use nu_engine::{eval_block, CallExt}; +use nu_protocol::ast::{Block, Call}; +use nu_protocol::engine::{CaptureBlock, Command, EngineState, Stack}; +use nu_protocol::{ + Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, + PipelineIterator, ShellError, Signature, Span, SyntaxShape, Value, +}; +use std::collections::HashSet; +use std::iter::FromIterator; + +#[derive(Clone)] +pub struct UpdateCells; + +impl Command for UpdateCells { + fn name(&self) -> &str { + "update cells" + } + + fn signature(&self) -> Signature { + Signature::build("update cells") + .required( + "block", + SyntaxShape::Block(Some(vec![SyntaxShape::Any])), + "the block to run an update for each cell", + ) + .named( + "columns", + SyntaxShape::Table, + "list of columns to update", + Some('c'), + ) + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Update the table cells." + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Update the zero value cells to empty strings.", + example: r#"[ + [2021-04-16, 2021-06-10, 2021-09-18, 2021-10-15, 2021-11-16, 2021-11-17, 2021-11-18]; + [ 37, 0, 0, 0, 37, 0, 0] +] | update cells {|value| + if $value == 0 { + "" + } else { + $value + } +}"#, + result: Some(Value::List { + vals: vec![Value::Record { + cols: vec![ + "2021-04-16".into(), + "2021-06-10".into(), + "2021-09-18".into(), + "2021-10-15".into(), + "2021-11-16".into(), + "2021-11-17".into(), + "2021-11-18".into(), + ], + vals: vec![ + Value::test_int(37), + Value::test_string(""), + Value::test_string(""), + Value::test_string(""), + Value::test_int(37), + Value::test_string(""), + Value::test_string(""), + ], + span: Span::test_data(), + }], + span: Span::test_data(), + }), + }, + Example { + description: "Update the zero value cells to empty strings in 2 last columns.", + example: r#"[ + [2021-04-16, 2021-06-10, 2021-09-18, 2021-10-15, 2021-11-16, 2021-11-17, 2021-11-18]; + [ 37, 0, 0, 0, 37, 0, 0] +] | update cells -c ["2021-11-18", "2021-11-17"] {|value| + if $value == 0 { + "" + } else { + $value + } +}"#, + result: Some(Value::List { + vals: vec![Value::Record { + cols: vec![ + "2021-04-16".into(), + "2021-06-10".into(), + "2021-09-18".into(), + "2021-10-15".into(), + "2021-11-16".into(), + "2021-11-17".into(), + "2021-11-18".into(), + ], + vals: vec![ + Value::test_int(37), + Value::test_int(0), + Value::test_int(0), + Value::test_int(0), + Value::test_int(37), + Value::test_string(""), + Value::test_string(""), + ], + span: Span::test_data(), + }], + span: Span::test_data(), + }), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + // the block to run on each cell + let engine_state = engine_state.clone(); + let block: CaptureBlock = call.req(&engine_state, stack, 0)?; + let mut stack = stack.captures_to_stack(&block.captures); + let orig_env_vars = stack.env_vars.clone(); + let orig_env_hidden = stack.env_hidden.clone(); + + let ctrlc = engine_state.ctrlc.clone(); + let block: Block = engine_state.get_block(block.block_id).clone(); + + let span = call.head; + + stack.with_env(&orig_env_vars, &orig_env_hidden); + + // the columns to update + let columns: Option = call.get_flag(&engine_state, &mut stack, "columns")?; + let columns: Option> = match columns { + Some(val) => { + let cols = val + .as_list()? + .iter() + .map(|val| val.as_string()) + .collect::, ShellError>>()?; + Some(HashSet::from_iter(cols.into_iter())) + } + None => None, + }; + + Ok(UpdateCellIterator { + input: input.into_iter(), + engine_state, + stack, + block, + columns, + span, + } + .into_pipeline_data(ctrlc)) + } +} + +struct UpdateCellIterator { + input: PipelineIterator, + columns: Option>, + engine_state: EngineState, + stack: Stack, + block: Block, + span: Span, +} + +impl Iterator for UpdateCellIterator { + type Item = Value; + + fn next(&mut self) -> Option { + match self.input.next() { + Some(val) => { + if let Some(ref cols) = self.columns { + if !val.columns().iter().any(|c| cols.contains(c)) { + return Some(val); + } + } + + match val { + Value::Record { vals, cols, span } => Some(Value::Record { + vals: cols + .iter() + .zip(vals.into_iter()) + .map(|(col, val)| match &self.columns { + Some(cols) if !cols.contains(col) => val, + _ => process_cell( + val, + &self.engine_state, + &mut self.stack, + &self.block, + span, + ), + }) + .collect(), + cols, + span, + }), + val => Some(process_cell( + val, + &self.engine_state, + &mut self.stack, + &self.block, + self.span, + )), + } + } + None => None, + } + } +} + +fn process_cell( + val: Value, + engine_state: &EngineState, + stack: &mut Stack, + block: &Block, + span: Span, +) -> Value { + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + stack.add_var(*var_id, val.clone()); + } + } + match eval_block(engine_state, stack, block, val.into_pipeline_data()) { + Ok(pd) => pd.into_value(span), + Err(e) => Value::Error { error: e }, + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(UpdateCells {}) + } +}