From 3c8716873ee54ad20740be5d270494fb8d89a706 Mon Sep 17 00:00:00 2001 From: Stefan Stanciulescu <71919805+onthebridgetonowhere@users.noreply.github.com> Date: Sat, 29 Jan 2022 21:47:28 +0100 Subject: [PATCH] Port rotate (#880) * Add rotate command * Add rotate counter clockwise * Fix comments in the code * Fix clippy warnings * Fix comment * Fix wrong step for non even table sizes * Fix comment for moving through array * Refactor rotate and have only one command with a --ccw flag for counter-clockwise rotation. By default, rotate is clockwise * Update usage description --- crates/nu-command/src/default_context.rs | 1 + crates/nu-command/src/filters/mod.rs | 2 + crates/nu-command/src/filters/rotate.rs | 360 +++++++++++++++++++++++ 3 files changed, 363 insertions(+) create mode 100644 crates/nu-command/src/filters/rotate.rs diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index a1380f9cfc..676cce4468 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -87,6 +87,7 @@ pub fn create_default_context(cwd: impl AsRef) -> EngineState { Reject, Rename, Reverse, + Rotate, Select, Shuffle, Skip, diff --git a/crates/nu-command/src/filters/mod.rs b/crates/nu-command/src/filters/mod.rs index 4b5e73bd93..9da0d7ab50 100644 --- a/crates/nu-command/src/filters/mod.rs +++ b/crates/nu-command/src/filters/mod.rs @@ -28,6 +28,7 @@ mod reduce; mod reject; mod rename; mod reverse; +mod rotate; mod select; mod shuffle; mod skip; @@ -69,6 +70,7 @@ pub use reduce::Reduce; pub use reject::Reject; pub use rename::Rename; pub use reverse::Reverse; +pub use rotate::Rotate; pub use select::Select; pub use shuffle::Shuffle; pub use skip::*; diff --git a/crates/nu-command/src/filters/rotate.rs b/crates/nu-command/src/filters/rotate.rs new file mode 100644 index 0000000000..ace6e927cd --- /dev/null +++ b/crates/nu-command/src/filters/rotate.rs @@ -0,0 +1,360 @@ +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, SyntaxShape, + Value, +}; + +#[derive(Clone)] +pub struct Rotate; + +impl Command for Rotate { + fn name(&self) -> &str { + "rotate" + } + + fn signature(&self) -> Signature { + Signature::build("rotate") + .switch("ccw", "rotate counter clockwise", None) + .rest( + "rest", + SyntaxShape::String, + "the names to give columns once rotated", + ) + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Rotates a table clockwise (default) or counter-clockwise (use --ccw flag)." + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Rotate 2x2 table clockwise", + example: "[[a b]; [1 2]] | rotate", + result: Some(Value::List { + vals: vec![ + Value::Record { + cols: vec!["Column0".to_string(), "Column1".to_string()], + vals: vec![Value::test_int(1), Value::test_string("a")], + span: Span::test_data(), + }, + Value::Record { + cols: vec!["Column0".to_string(), "Column1".to_string()], + vals: vec![Value::test_int(2), Value::test_string("b")], + span: Span::test_data(), + }, + ], + span: Span::test_data(), + }), + }, + Example { + description: "Rotate 2x3 table clockwise", + example: "[[a b]; [1 2] [3 4] [5 6]] | rotate", + result: Some(Value::List { + vals: vec![ + Value::Record { + cols: vec![ + "Column0".to_string(), + "Column1".to_string(), + "Column2".to_string(), + "Column3".to_string(), + ], + vals: vec![ + Value::test_int(5), + Value::test_int(3), + Value::test_int(1), + Value::test_string("a"), + ], + span: Span::test_data(), + }, + Value::Record { + cols: vec![ + "Column0".to_string(), + "Column1".to_string(), + "Column2".to_string(), + "Column3".to_string(), + ], + vals: vec![ + Value::test_int(6), + Value::test_int(4), + Value::test_int(2), + Value::test_string("b"), + ], + span: Span::test_data(), + }, + ], + span: Span::test_data(), + }), + }, + Example { + description: "Rotate table clockwise and change columns names", + example: "[[a b]; [1 2]] | rotate col_a col_b", + result: Some(Value::List { + vals: vec![ + Value::Record { + cols: vec!["col_a".to_string(), "col_b".to_string()], + vals: vec![Value::test_int(1), Value::test_string("a")], + span: Span::test_data(), + }, + Value::Record { + cols: vec!["col_a".to_string(), "col_b".to_string()], + vals: vec![Value::test_int(2), Value::test_string("b")], + span: Span::test_data(), + }, + ], + span: Span::test_data(), + }), + }, + Example { + description: "Rotate table counter clockwise", + example: "[[a b]; [1 2]] | rotate --ccw", + result: Some(Value::List { + vals: vec![ + Value::Record { + cols: vec!["Column0".to_string(), "Column1".to_string()], + vals: vec![Value::test_string("b"), Value::test_int(2)], + span: Span::test_data(), + }, + Value::Record { + cols: vec!["Column0".to_string(), "Column1".to_string()], + vals: vec![Value::test_string("a"), Value::test_int(1)], + span: Span::test_data(), + }, + ], + span: Span::test_data(), + }), + }, + Example { + description: "Rotate table counter-clockwise", + example: "[[a b]; [1 2] [3 4] [5 6]] | rotate --ccw", + result: Some(Value::List { + vals: vec![ + Value::Record { + cols: vec![ + "Column0".to_string(), + "Column1".to_string(), + "Column2".to_string(), + "Column3".to_string(), + ], + vals: vec![ + Value::test_string("b"), + Value::test_int(2), + Value::test_int(4), + Value::test_int(6), + ], + span: Span::test_data(), + }, + Value::Record { + cols: vec![ + "Column0".to_string(), + "Column1".to_string(), + "Column2".to_string(), + "Column3".to_string(), + ], + vals: vec![ + Value::test_string("a"), + Value::test_int(1), + Value::test_int(3), + Value::test_int(5), + ], + span: Span::test_data(), + }, + ], + span: Span::test_data(), + }), + }, + Example { + description: "Rotate table counter-clockwise and change columns names", + example: "[[a b]; [1 2]] | rotate --ccw col_a col_b", + result: Some(Value::List { + vals: vec![ + Value::Record { + cols: vec!["col_a".to_string(), "col_b".to_string()], + vals: vec![Value::test_string("b"), Value::test_int(2)], + span: Span::test_data(), + }, + Value::Record { + cols: vec!["col_a".to_string(), "col_b".to_string()], + vals: vec![Value::test_string("a"), Value::test_int(1)], + span: Span::test_data(), + }, + ], + span: Span::test_data(), + }), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + rotate(engine_state, stack, call, input) + } +} + +pub fn rotate( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let col_given_names: Vec = call.rest(engine_state, stack, 0)?; + let mut values = input.into_iter().collect::>(); + let mut old_column_names = vec![]; + let mut new_values = vec![]; + let mut not_a_record = false; + let total_rows = &mut values.len(); + let ccw: bool = call.has_flag("ccw"); + + if !ccw { + values.reverse(); + } + + if !values.is_empty() { + for val in values.into_iter() { + match val { + Value::Record { + cols, + vals, + span: _, + } => { + old_column_names = cols; + for v in vals { + new_values.push(v) + } + } + Value::List { vals, span: _ } => { + not_a_record = true; + for v in vals { + new_values.push(v); + } + } + Value::String { val, span } => { + not_a_record = true; + new_values.push(Value::String { val, span }) + } + x => { + not_a_record = true; + new_values.push(x) + } + } + } + } else { + return Err(ShellError::UnsupportedInput( + "Rotate command requires a Nu value as input".to_string(), + call.head, + )); + } + + let total_columns = &old_column_names.len(); + + // we use this for building columns names, but for non-records we get an extra row so we remove it + if *total_columns == 0 { + *total_rows -= 1; + } + + // holder for the new column names, particularly if none are provided by the user we create names as Column0, Column1, etc. + let mut new_column_names = { + let mut res = vec![]; + for idx in 0..(*total_rows + 1) { + res.push(format!("Column{}", idx)); + } + res.to_vec() + }; + + // we got new names for columns from the input, so we need to swap those we already made + if !col_given_names.is_empty() { + for (idx, val) in col_given_names.into_iter().enumerate() { + if idx > new_column_names.len() - 1 { + break; + } + new_column_names[idx] = val; + } + } + + if not_a_record { + return Ok(Value::List { + vals: vec![Value::Record { + cols: new_column_names, + vals: new_values, + span: call.head, + }], + span: call.head, + } + .into_pipeline_data()); + } + + // holder for the new records + let mut final_values = vec![]; + + // the number of initial columns will be our number of rows, so we iterate through that to get the new number of rows that we need to make + // for counter clockwise, we're iterating from right to left and have a pair of (index, value) + let columns_iter = if ccw { + old_column_names + .iter() + .enumerate() + .rev() + .collect::>() + } else { + // as we're rotating clockwise, we're iterating from left to right and have a pair of (index, value) + old_column_names.iter().enumerate().collect::>() + }; + + for (idx, val) in columns_iter { + // when rotating counter clockwise, the old columns names become the first column's values + let mut res = if ccw { + vec![Value::String { + val: val.to_string(), + span: call.head, + }] + } else { + vec![] + }; + + let new_vals = { + // move through the array with a step, which is every new_values size / total rows, starting from our old column's index + // so if initial data was like this [[a b]; [1 2] [3 4]] - we basically iterate on this [3 4 1 2] array, so we pick 3, then 1, and then when idx increases, we pick 4 and 2 + for i in (idx..new_values.len()).step_by(new_values.len() / *total_rows) { + res.push(new_values[i].clone()); + } + // when rotating clockwise, the old column names become the last column's values + if !ccw { + res.push(Value::String { + val: val.to_string(), + span: call.head, + }); + } + res.to_vec() + }; + final_values.push(Value::Record { + cols: new_column_names.clone(), + vals: new_vals, + span: call.head, + }) + } + + Ok(Value::List { + vals: final_values, + span: call.head, + } + .into_pipeline_data()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(Rotate) + } +}