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
This commit is contained in:
Stefan Stanciulescu 2022-01-29 21:47:28 +01:00 committed by GitHub
parent 44821d9941
commit 3c8716873e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 363 additions and 0 deletions

View file

@ -87,6 +87,7 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
Reject,
Rename,
Reverse,
Rotate,
Select,
Shuffle,
Skip,

View file

@ -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::*;

View file

@ -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<Example> {
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<PipelineData, ShellError> {
rotate(engine_state, stack, call, input)
}
}
pub fn rotate(
engine_state: &EngineState,
stack: &mut Stack,
call: &Call,
input: PipelineData,
) -> Result<PipelineData, ShellError> {
let col_given_names: Vec<String> = call.rest(engine_state, stack, 0)?;
let mut values = input.into_iter().collect::<Vec<_>>();
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::<Vec<_>>()
} 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::<Vec<_>>()
};
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)
}
}