Merge pull request #82 from nushell/add_select

Add select

crates/nu-command/src/default_context.rs

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

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;

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<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));
+//     }
+// }

crates/nu-protocol/src/ast/cell_path.rs

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

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

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