Add select

This commit is contained in:
JT 2021-10-02 17:55:05 +13:00
parent 7899745c4b
commit 6b76dd7cd7
6 changed files with 213 additions and 2 deletions

View File

@ -33,6 +33,7 @@ pub fn create_default_context() -> Rc<RefCell<EngineState>> {
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));

View File

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

View File

@ -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<nu_protocol::Value, nu_protocol::ShellError> {
let columns: Vec<CellPath> = call.rest(context, 0)?;
let span = call.head;
select(span, columns, input)
}
fn examples(&self) -> Vec<Example> {
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<CellPath>, input: Value) -> Result<Value, ShellError> {
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));
// }
// }

View File

@ -13,6 +13,24 @@ pub struct CellPath {
pub members: Vec<PathMember>,
}
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,

View File

@ -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(),
}
}

View File

@ -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",
)
}