From 1cbb785969d665d3a2ec7c54c0d763d186482fe6 Mon Sep 17 00:00:00 2001 From: Michael Angerman Date: Tue, 14 Dec 2021 11:54:27 -0800 Subject: [PATCH] port over from nushell drop column (#495) * port over from nushell drop column * fix clippy --- crates/nu-command/src/default_context.rs | 1 + crates/nu-command/src/filters/drop/column.rs | 163 +++++++++++++++++++ crates/nu-command/src/filters/drop/mod.rs | 2 + src/tests.rs | 12 ++ 4 files changed, 178 insertions(+) create mode 100644 crates/nu-command/src/filters/drop/column.rs diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 07932b95fa..8743082ce2 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -50,6 +50,7 @@ pub fn create_default_context() -> EngineState { Append, Collect, Drop, + DropColumn, Each, First, Get, diff --git a/crates/nu-command/src/filters/drop/column.rs b/crates/nu-command/src/filters/drop/column.rs new file mode 100644 index 0000000000..ae99ad0c19 --- /dev/null +++ b/crates/nu-command/src/filters/drop/column.rs @@ -0,0 +1,163 @@ +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 DropColumn; + +impl Command for DropColumn { + fn name(&self) -> &str { + "drop column" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .optional( + "columns", + SyntaxShape::Int, + "starting from the end, the number of columns to remove", + ) + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Remove the last number of columns. If you want to remove columns by name, try 'reject'." + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + // the number of columns to drop + let columns: Option = call.opt(engine_state, stack, 0)?; + let span = call.head; + + let columns_to_drop = if let Some(quantity) = columns { + quantity + } else { + 1 + }; + + dropcol(engine_state, span, input, columns_to_drop) + } +} + +fn dropcol( + engine_state: &EngineState, + span: Span, + input: PipelineData, + columns: i64, // the number of columns to drop +) -> Result { + 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(input: Vec, mut num_of_columns_to_drop: i64) -> Vec { + let vlen: i64 = input.len() as i64; + + if num_of_columns_to_drop > vlen { + num_of_columns_to_drop = vlen; + } + + let num_of_columns_to_keep = (vlen - num_of_columns_to_drop) as usize; + input[0..num_of_columns_to_keep].to_vec() +} diff --git a/crates/nu-command/src/filters/drop/mod.rs b/crates/nu-command/src/filters/drop/mod.rs index 70dd1f6fa0..9b9da9cebf 100644 --- a/crates/nu-command/src/filters/drop/mod.rs +++ b/crates/nu-command/src/filters/drop/mod.rs @@ -1,3 +1,5 @@ +pub mod column; pub mod command; +pub use column::DropColumn; pub use command::Drop; diff --git a/src/tests.rs b/src/tests.rs index 17b8d274ea..1077bd811a 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1233,6 +1233,18 @@ fn command_filter_reject_3() -> TestResult { ) } +#[test] +fn command_drop_column_1() -> TestResult { + run_test( + "[[lang, gems, grade]; [nu, 100, a]] | drop column 2 | to json", + r#"[ + { + "lang": "nu" + } +]"#, + ) +} + #[test] fn chained_operator_typecheck() -> TestResult { run_test("1 != 2 && 3 != 4 && 5 != 6", "true")