diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 818d278ff..dc119edc5 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -33,6 +33,7 @@ pub fn create_default_context() -> Rc> { working_set.add_decl(Box::new(Ls)); working_set.add_decl(Box::new(Module)); working_set.add_decl(Box::new(Ps)); + working_set.add_decl(Box::new(Select)); working_set.add_decl(Box::new(Sys)); working_set.add_decl(Box::new(Table)); working_set.add_decl(Box::new(Use)); diff --git a/crates/nu-command/src/filters/mod.rs b/crates/nu-command/src/filters/mod.rs index ba3508fd6..c90572156 100644 --- a/crates/nu-command/src/filters/mod.rs +++ b/crates/nu-command/src/filters/mod.rs @@ -3,6 +3,7 @@ mod for_; mod get; mod length; mod lines; +mod select; mod where_; mod wrap; @@ -11,5 +12,6 @@ pub use for_::For; pub use get::Get; pub use length::Length; pub use lines::Lines; +pub use select::Select; pub use where_::Where; pub use wrap::Wrap; diff --git a/crates/nu-command/src/filters/select.rs b/crates/nu-command/src/filters/select.rs new file mode 100644 index 000000000..8611517e4 --- /dev/null +++ b/crates/nu-command/src/filters/select.rs @@ -0,0 +1,182 @@ +use nu_engine::CallExt; +use nu_protocol::ast::{Call, CellPath}; +use nu_protocol::engine::{Command, EvaluationContext}; +use nu_protocol::{Example, IntoValueStream, ShellError, Signature, Span, SyntaxShape, Value}; + +pub struct Select; + +impl Command for Select { + fn name(&self) -> &str { + "select" + } + + fn signature(&self) -> Signature { + Signature::build("select").rest( + "rest", + SyntaxShape::CellPath, + "the columns to select from the table", + ) + } + + fn usage(&self) -> &str { + "Down-select table to only these columns." + } + + fn run( + &self, + context: &EvaluationContext, + call: &Call, + input: Value, + ) -> Result { + let columns: Vec = call.rest(context, 0)?; + let span = call.head; + + select(span, columns, input) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Select just the name column", + example: "ls | select name", + result: None, + }, + Example { + description: "Select the name and size columns", + example: "ls | select name size", + result: None, + }, + ] + } +} + +fn select(span: Span, columns: Vec, input: Value) -> Result { + if columns.is_empty() { + return Err(ShellError::CantFindColumn(span)); + } + + match input { + Value::List { + vals: input_vals, + span, + } => { + let mut output = vec![]; + + for input_val in input_vals { + let mut cols = vec![]; + let mut vals = vec![]; + for path in &columns { + //FIXME: improve implementation to not clone + 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(Value::List { vals: output, span }) + } + Value::Stream { stream, span } => Ok(Value::Stream { + stream: stream + .map(move |x| { + let mut cols = vec![]; + let mut vals = vec![]; + for path in &columns { + //FIXME: improve implementation to not clone + match x.clone().follow_cell_path(&path.members) { + Ok(value) => { + cols.push(path.into_string()); + vals.push(value); + } + Err(error) => { + cols.push(path.into_string()); + vals.push(Value::Error { error }); + } + } + } + + Value::Record { cols, vals, span } + }) + .into_value_stream(), + span, + }), + v => { + let mut cols = vec![]; + let mut vals = vec![]; + + for cell_path in columns { + // FIXME: remove clone + 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 }) + } + } +} + +// #[cfg(test)] +// mod tests { +// use nu_protocol::ColumnPath; +// use nu_source::Span; +// use nu_source::SpannedItem; +// use nu_source::Tag; +// use nu_stream::InputStream; +// use nu_test_support::value::nothing; +// use nu_test_support::value::row; +// use nu_test_support::value::string; + +// use super::select; +// use super::Command; +// use super::ShellError; + +// #[test] +// fn examples_work_as_expected() -> Result<(), ShellError> { +// use crate::examples::test as test_examples; + +// test_examples(Command {}) +// } + +// #[test] +// fn select_using_sparse_table() { +// // Create a sparse table with 3 rows: +// // col_foo | col_bar +// // ----------------- +// // foo | +// // | bar +// // foo | +// let input = vec![ +// row(indexmap! {"col_foo".into() => string("foo")}), +// row(indexmap! {"col_bar".into() => string("bar")}), +// row(indexmap! {"col_foo".into() => string("foo")}), +// ]; + +// let expected = vec![ +// row( +// indexmap! {"col_none".into() => nothing(), "col_foo".into() => string("foo"), "col_bar".into() => nothing()}, +// ), +// row( +// indexmap! {"col_none".into() => nothing(), "col_foo".into() => nothing(), "col_bar".into() => string("bar")}, +// ), +// row( +// indexmap! {"col_none".into() => nothing(), "col_foo".into() => string("foo"), "col_bar".into() => nothing()}, +// ), +// ]; + +// let actual = select( +// Tag::unknown(), +// vec![ +// ColumnPath::build(&"col_none".to_string().spanned(Span::unknown())), +// ColumnPath::build(&"col_foo".to_string().spanned(Span::unknown())), +// ColumnPath::build(&"col_bar".to_string().spanned(Span::unknown())), +// ], +// input.into(), +// ); + +// assert_eq!(Ok(expected), actual.map(InputStream::into_vec)); +// } +// } diff --git a/crates/nu-protocol/src/ast/cell_path.rs b/crates/nu-protocol/src/ast/cell_path.rs index b6071b644..078b64919 100644 --- a/crates/nu-protocol/src/ast/cell_path.rs +++ b/crates/nu-protocol/src/ast/cell_path.rs @@ -13,6 +13,24 @@ pub struct CellPath { pub members: Vec, } +impl CellPath { + pub fn into_string(&self) -> String { + let mut output = String::new(); + + for (idx, elem) in self.members.iter().enumerate() { + if idx > 0 { + output.push('.'); + } + match elem { + PathMember::Int { val, .. } => output.push_str(&format!("{}", val)), + PathMember::String { val, .. } => output.push_str(val), + } + } + + output + } +} + #[derive(Debug, Clone)] pub struct FullCellPath { pub head: Expression, diff --git a/crates/nu-protocol/src/value/mod.rs b/crates/nu-protocol/src/value/mod.rs index aed1aa832..eae64b8dc 100644 --- a/crates/nu-protocol/src/value/mod.rs +++ b/crates/nu-protocol/src/value/mod.rs @@ -191,7 +191,7 @@ impl Value { Value::Nothing { .. } => String::new(), Value::Error { error } => format!("{:?}", error), Value::Binary { val, .. } => format!("{:?}", val), - Value::CellPath { val, .. } => format!("{:?}", val), + Value::CellPath { val, .. } => val.into_string(), } } @@ -223,7 +223,7 @@ impl Value { Value::Nothing { .. } => String::new(), Value::Error { error } => format!("{:?}", error), Value::Binary { val, .. } => format!("{:?}", val), - Value::CellPath { val, .. } => format!("{:?}", val), + Value::CellPath { .. } => self.into_string(), } } diff --git a/src/tests.rs b/src/tests.rs index 37adcb609..914f0da51 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -409,3 +409,11 @@ fn get() -> TestResult { "B", ) } + +#[test] +fn select() -> TestResult { + run_test( + r#"([[name, age]; [a, 1], [b, 2]]) | select name | get 1 | get name"#, + "b", + ) +}