From 3ed45c7ba856ae2678d5233bcd9039258c1558f7 Mon Sep 17 00:00:00 2001
From: Darren Schroeder <343840+fdncred@users.noreply.github.com>
Date: Tue, 15 Aug 2023 07:01:45 -0500
Subject: [PATCH] allow `select` to take a $variable with a list of columns
 (#9987)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# Description

This PR enables `select` to take a constructed list of columns as a
variable.

```nushell
> let cols = [name type];[[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb]] | select $cols
  ╭#┬───name───┬type╮
  │0│Cargo.toml│toml│
  │1│Cargo.lock│toml│
  ╰─┴──────────┴────╯
```
and rows
```nushell
> let rows = [0 2];[[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb] [file.json json 3kb]] | select $rows
  ╭#┬───name───┬type┬size╮
  │0│Cargo.toml│toml│1kb │
  │1│file.json │json│3kb │
  ╰─┴──────────┴────┴────╯
```

# User-Facing Changes
<!-- List of all changes that impact the user experience here. This
helps us keep track of breaking changes. -->

# Tests + Formatting
<!--
Don't forget to add tests that cover your changes.

Make sure you've run and fixed any issues with these commands:

- `cargo fmt --all -- --check` to check standard code formatting (`cargo
fmt --all` applies these changes)
- `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used -A
clippy::needless_collect -A clippy::result_large_err` to check that
you're using the standard code style
- `cargo test --workspace` to check that all tests pass
- `cargo run -- -c "use std testing; testing run-tests --path
crates/nu-std"` to run the tests for the standard library

> **Note**
> from `nushell` you can also use the `toolkit` as follows
> ```bash
> use toolkit.nu # or use an `env_change` hook to activate it
automatically
> toolkit check pr
> ```
-->

# After Submitting
<!-- If your PR had any user-facing changes, update [the
documentation](https://github.com/nushell/nushell.github.io) after the
PR is merged, if necessary. This will help us keep the docs up to date.
-->
---
 crates/nu-command/src/filters/select.rs    | 81 ++++++++++++++++++++--
 crates/nu-command/tests/commands/select.rs | 20 ++++++
 2 files changed, 97 insertions(+), 4 deletions(-)

diff --git a/crates/nu-command/src/filters/select.rs b/crates/nu-command/src/filters/select.rs
index 0828c3d209..e8396a82e5 100644
--- a/crates/nu-command/src/filters/select.rs
+++ b/crates/nu-command/src/filters/select.rs
@@ -30,7 +30,10 @@ impl Command for Select {
             )
             .rest(
                 "rest",
-                SyntaxShape::CellPath,
+                SyntaxShape::OneOf(vec![
+                    SyntaxShape::CellPath,
+                    SyntaxShape::List(Box::new(SyntaxShape::CellPath)),
+                ]),
                 "the columns to select from the table",
             )
             .allow_variants_without_examples(true)
@@ -58,17 +61,77 @@ produce a table, a list will produce a list, and a record will produce a record.
         call: &Call,
         input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
-        let mut columns: Vec<CellPath> = call.rest(engine_state, stack, 0)?;
+        let columns: Vec<Value> = call.rest(engine_state, stack, 0)?;
+        let mut new_columns: Vec<CellPath> = vec![];
+        for col_val in columns {
+            match col_val {
+                Value::CellPath { val, .. } => {
+                    new_columns.push(val);
+                }
+                Value::List { vals, .. } => {
+                    for val in vals {
+                        match val {
+                            Value::String { val, .. } => {
+                                let cv = CellPath {
+                                    members: vec![PathMember::String {
+                                        val: val.clone(),
+                                        span: Span::unknown(),
+                                        optional: false,
+                                    }],
+                                };
+                                new_columns.push(cv.clone());
+                            }
+                            Value::Int { val, .. } => {
+                                let cv = CellPath {
+                                    members: vec![PathMember::Int {
+                                        val: val as usize,
+                                        span: Span::unknown(),
+                                        optional: false,
+                                    }],
+                                };
+                                new_columns.push(cv.clone());
+                            }
+                            y => {
+                                return Err(ShellError::CantConvert {
+                                    to_type: "cell path".into(),
+                                    from_type: y.get_type().to_string(),
+                                    span: y.span()?,
+                                    help: None,
+                                });
+                            }
+                        }
+                    }
+                }
+                Value::String { val, .. } => {
+                    let cv = CellPath {
+                        members: vec![PathMember::String {
+                            val: val.clone(),
+                            span: Span::unknown(),
+                            optional: false,
+                        }],
+                    };
+                    new_columns.push(cv.clone());
+                }
+                x => {
+                    return Err(ShellError::CantConvert {
+                        to_type: "cell path".into(),
+                        from_type: x.get_type().to_string(),
+                        span: x.span()?,
+                        help: None,
+                    });
+                }
+            }
+        }
         let ignore_errors = call.has_flag("ignore-errors");
         let span = call.head;
 
         if ignore_errors {
-            for cell_path in &mut columns {
+            for cell_path in &mut new_columns {
                 cell_path.make_optional();
             }
         }
 
-        select(engine_state, span, columns, input)
+        select(engine_state, span, new_columns, input)
     }
 
     fn examples(&self) -> Vec<Example> {
@@ -96,6 +159,16 @@ produce a table, a list will produce a list, and a record will produce a record.
                 example: "ls | select 0 1 2 3",
                 result: None,
             },
+            Example {
+                description: "Select columns by a provided list of columns",
+                example: "let cols = [name type];[[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb]] | select $cols",
+                result: None
+            },
+            Example {
+                description: "Select rows by a provided list of rows",
+                example: "let rows = [0 2];[[name type size]; [Cargo.toml toml 1kb] [Cargo.lock toml 2kb] [file.json json 3kb]] | select $rows",
+                result: None
+            },
         ]
     }
 }
diff --git a/crates/nu-command/tests/commands/select.rs b/crates/nu-command/tests/commands/select.rs
index f761628f8c..665a2fd583 100644
--- a/crates/nu-command/tests/commands/select.rs
+++ b/crates/nu-command/tests/commands/select.rs
@@ -244,3 +244,23 @@ fn select_on_empty_list_returns_empty_list() {
     let actual = nu!("[] | each {|i| $i} | select foo | to nuon");
     assert_eq!(actual.out, "[]");
 }
+
+#[test]
+fn select_columns_with_variable_list() {
+    let actual = nu!(r#"
+        let columns = [a c];
+        echo [[a b c]; [1 2 3]] | select $columns | to nuon
+        "#);
+
+    assert_eq!(actual.out, "[[a, c]; [1, 3]]");
+}
+
+#[test]
+fn select_rows_with_variable_list() {
+    let actual = nu!(r#"
+        let rows = [0 2];
+        echo [[a b c]; [1 2 3] [4 5 6] [7 8 9]] | select $rows | to nuon
+        "#);
+
+    assert_eq!(actual.out, "[[a, b, c]; [1, 2, 3], [7, 8, 9]]");
+}