From 574c5961c847f8ca06d80a4d7a148a751556d2eb Mon Sep 17 00:00:00 2001
From: Josh Cheek <josh.cheek@gmail.com>
Date: Mon, 4 Oct 2021 19:23:37 -0500
Subject: [PATCH] Add `-c` flag to `select` command (#4062)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

See cc3653cfd981672a95bb3942dbdc718290bc5453 for more on the `-c` flag.

Co-authored-by: Andrés N. Robalino <andres@androbtech.com>

Co-authored-by: Andrés N. Robalino <andres@androbtech.com>
---
 .../nu-command/src/commands/filters/select.rs | 62 ++++++++++++++++---
 1 file changed, 55 insertions(+), 7 deletions(-)

diff --git a/crates/nu-command/src/commands/filters/select.rs b/crates/nu-command/src/commands/filters/select.rs
index f098a88921..41e1ff1261 100644
--- a/crates/nu-command/src/commands/filters/select.rs
+++ b/crates/nu-command/src/commands/filters/select.rs
@@ -15,11 +15,18 @@ impl WholeStreamCommand for Command {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build("select").rest(
-            "rest",
-            SyntaxShape::ColumnPath,
-            "the columns to select from the table",
-        )
+        Signature::build("select")
+            .named(
+                "columns",
+                SyntaxShape::Table,
+                "Optionally operate by column path",
+                Some('c'),
+            )
+            .rest(
+                "rest",
+                SyntaxShape::ColumnPath,
+                "the columns to select from the table",
+            )
     }
 
     fn usage(&self) -> &str {
@@ -27,10 +34,10 @@ impl WholeStreamCommand for Command {
     }
 
     fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
-        let columns: Vec<ColumnPath> = args.rest(0)?;
+        let mut columns = args.rest(0)?;
+        columns.extend(column_paths_from_args(&args)?);
         let input = args.input;
         let name = args.call_info.name_tag;
-
         select(name, columns, input)
     }
 
@@ -46,10 +53,51 @@ impl WholeStreamCommand for Command {
                 example: "ls | select name size",
                 result: None,
             },
+            Example {
+                description: "Select columns dynamically",
+                example: "[[a b]; [1 2]] | select -c [a]",
+                result: Some(vec![UntaggedValue::row(indexmap! {
+                    "a".to_string() => UntaggedValue::int(1).into(),
+                })
+                .into()]),
+            },
         ]
     }
 }
 
+fn column_paths_from_args(args: &CommandArgs) -> Result<Vec<ColumnPath>, ShellError> {
+    let column_paths: Option<Vec<Value>> = args.get_flag("columns")?;
+    let has_columns = column_paths.is_some();
+    let column_paths = match column_paths {
+        Some(cols) => {
+            let mut c = Vec::new();
+            for col in cols {
+                let colpath = ColumnPath::build(&col.convert_to_string().spanned_unknown());
+                if !colpath.is_empty() {
+                    c.push(colpath)
+                }
+            }
+            c
+        }
+        None => Vec::new(),
+    };
+
+    if has_columns && column_paths.is_empty() {
+        let colval: Option<Value> = args.get_flag("columns")?;
+        let colspan = match colval {
+            Some(v) => v.tag.span,
+            None => Span::unknown(),
+        };
+        return Err(ShellError::labeled_error(
+            "Requires a list of columns",
+            "must be a list of columns",
+            colspan,
+        ));
+    }
+
+    Ok(column_paths)
+}
+
 fn select(
     name: Tag,
     columns: Vec<ColumnPath>,