From 03e22b071ac7cf211c83cf89925fb6261fa098ca Mon Sep 17 00:00:00 2001 From: Michael Angerman Date: Sat, 4 Dec 2021 19:09:45 -0800 Subject: [PATCH] port over the reject command from nushell (#419) * port over reject * add some tests to src/tests --- crates/nu-command/src/default_context.rs | 1 + crates/nu-command/src/filters/mod.rs | 2 + crates/nu-command/src/filters/reject.rs | 157 +++++++++++++++++++++++ src/tests.rs | 21 +++ 4 files changed, 181 insertions(+) create mode 100644 crates/nu-command/src/filters/reject.rs diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 3a3ca2470d..c43c0f5048 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -118,6 +118,7 @@ pub fn create_default_context() -> EngineState { Ps, Range, Random, + Reject, Reverse, Rm, Select, diff --git a/crates/nu-command/src/filters/mod.rs b/crates/nu-command/src/filters/mod.rs index 61f4f15939..9805e0809d 100644 --- a/crates/nu-command/src/filters/mod.rs +++ b/crates/nu-command/src/filters/mod.rs @@ -11,6 +11,7 @@ mod length; mod lines; mod par_each; mod range; +mod reject; mod reverse; mod select; mod shuffle; @@ -33,6 +34,7 @@ pub use length::Length; pub use lines::Lines; pub use par_each::ParEach; pub use range::Range; +pub use reject::Reject; pub use reverse::Reverse; pub use select::Select; pub use shuffle::Shuffle; diff --git a/crates/nu-command/src/filters/reject.rs b/crates/nu-command/src/filters/reject.rs new file mode 100644 index 0000000000..ccbb2915c9 --- /dev/null +++ b/crates/nu-command/src/filters/reject.rs @@ -0,0 +1,157 @@ +use nu_engine::CallExt; +use nu_protocol::ast::{Call, CellPath}; +use nu_protocol::engine::{Command, EngineState, Stack}; +use nu_protocol::{ + Category, FromValue, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, ShellError, + Signature, Span, SyntaxShape, Value, +}; + +#[derive(Clone)] +pub struct Reject; + +impl Command for Reject { + fn name(&self) -> &str { + "reject" + } + + fn signature(&self) -> Signature { + Signature::build("reject") + .rest( + "rest", + SyntaxShape::String, + "the names of columns to remove from the table", + ) + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Remove the given columns from the table. If you want to remove rows, try 'drop'." + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let columns: Vec = call.rest(engine_state, stack, 0)?; + let span = call.head; + reject(engine_state, span, input, columns) + } +} + +fn reject( + engine_state: &EngineState, + span: Span, + input: PipelineData, + columns: Vec, +) -> Result { + if columns.is_empty() { + return Err(ShellError::CantFindColumn(span, span)); + } + + let mut keep_columns = vec![]; + + match input { + PipelineData::Value( + Value::List { + vals: input_vals, + span, + }, + .., + ) => { + let mut output = vec![]; + let input_cols = get_input_cols(input_vals.clone()); + let kc = get_keep_columns(input_cols, columns); + keep_columns = get_cellpath_columns(kc); + + for input_val in input_vals { + let mut cols = vec![]; + let mut vals = vec![]; + + for path in &keep_columns { + let fetcher = input_val.clone().follow_cell_path(&path.members)?; + cols.push(path.into_string()); + vals.push(fetcher); + } + output.push(Value::Record { cols, vals, span }) + } + + Ok(output + .into_iter() + .into_pipeline_data(engine_state.ctrlc.clone())) + } + PipelineData::Stream(stream, ..) => { + let mut output = vec![]; + + let v: Vec<_> = stream.into_iter().collect(); + let input_cols = get_input_cols(v.clone()); + let kc = get_keep_columns(input_cols, columns); + keep_columns = get_cellpath_columns(kc); + + for input_val in v { + let mut cols = vec![]; + let mut vals = vec![]; + + for path in &keep_columns { + let fetcher = input_val.clone().follow_cell_path(&path.members)?; + cols.push(path.into_string()); + vals.push(fetcher); + } + output.push(Value::Record { cols, vals, span }) + } + + Ok(output + .into_iter() + .into_pipeline_data(engine_state.ctrlc.clone())) + } + PipelineData::Value(v, ..) => { + let mut cols = vec![]; + let mut vals = vec![]; + + for cell_path in &keep_columns { + let result = v.clone().follow_cell_path(&cell_path.members)?; + + cols.push(cell_path.into_string()); + vals.push(result); + } + + Ok(Value::Record { cols, vals, span }.into_pipeline_data()) + } + } +} + +fn get_input_cols(input: Vec) -> Vec { + let rec = input.first(); + match rec { + Some(Value::Record { cols, vals: _, .. }) => cols.to_vec(), + _ => vec!["".to_string()], + } +} + +fn get_cellpath_columns(keep_cols: Vec) -> Vec { + let mut output = vec![]; + for keep_col in keep_cols { + let span = Span::unknown(); + let val = Value::String { + val: keep_col, + span, + }; + let cell_path = match CellPath::from_value(&val) { + Ok(v) => v, + Err(_) => return vec![], + }; + output.push(cell_path); + } + output +} + +fn get_keep_columns(mut input: Vec, rejects: Vec) -> Vec { + for reject in rejects { + if let Some(index) = input.iter().position(|value| *value == reject) { + input.swap_remove(index); + } + } + input +} diff --git a/src/tests.rs b/src/tests.rs index 5b8a359a57..d54f5ffff6 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1187,3 +1187,24 @@ fn comment_skipping_2() -> TestResult { "40", ) } + +#[test] +fn command_filter_reject_1() -> TestResult { + run_test("[[lang, gems]; [nu, 100]] | reject gems", "{lang: nu}") +} + +#[test] +fn command_filter_reject_2() -> TestResult { + run_test( + "[[lang, gems, grade]; [nu, 100, a]] | reject gems grade", + "{lang: nu}", + ) +} + +#[test] +fn command_filter_reject_3() -> TestResult { + run_test( + "[[lang, gems, grade]; [nu, 100, a]] | reject grade gems", + "{lang: nu}", + ) +}