From e5405d7f5c92e6bb8d0e57f4fba7dd42e5392146 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20N=2E=20Robalino?= <andres@androbtech.com>
Date: Tue, 3 Dec 2019 02:26:01 -0500
Subject: [PATCH] embed as column when embedding a list

---
 src/plugins/embed.rs  | 48 +++++++++++++++++++++++++++++++++++--------
 tests/filters_test.rs | 34 ++++++++++++++++++++++++++++--
 2 files changed, 71 insertions(+), 11 deletions(-)

diff --git a/src/plugins/embed.rs b/src/plugins/embed.rs
index cce969867..0fd177a09 100644
--- a/src/plugins/embed.rs
+++ b/src/plugins/embed.rs
@@ -11,18 +11,40 @@ use nu_source::Tag;
 
 struct Embed {
     field: Option<String>,
+    are_all_rows: bool,
     values: Vec<Value>,
 }
 impl Embed {
     fn new() -> Embed {
         Embed {
             field: None,
+            are_all_rows: true,
             values: Vec::new(),
         }
     }
 
     fn embed(&mut self, value: Value) -> Result<(), ShellError> {
-        self.values.push(value);
+        match &value {
+            Value {
+                value: UntaggedValue::Row(_),
+                ..
+            } => {
+                self.values.push(value);
+            }
+            _ => {
+                self.are_all_rows = false;
+
+                self.values.push(
+                    value::row(indexmap! {
+                        match &self.field {
+                            Some(key) => key.clone(),
+                            None => "Column".into()
+                        } => value
+                    })
+                    .into_value(Tag::unknown()),
+                );
+            }
+        }
         Ok(())
     }
 }
@@ -58,15 +80,23 @@ impl Plugin for Embed {
     }
 
     fn end_filter(&mut self) -> Result<Vec<ReturnValue>, ShellError> {
-        let row = value::row(indexmap! {
-            match &self.field {
-                Some(key) => key.clone(),
-                None => "root".into(),
-            } => value::table(&self.values).into_value(Tag::unknown()),
-        })
-        .into_untagged_value();
+        if self.are_all_rows {
+            let row = value::row(indexmap! {
+                match &self.field {
+                    Some(key) => key.clone(),
+                    None => "Column".into(),
+                } => value::table(&self.values).into_value(Tag::unknown()),
+            })
+            .into_untagged_value();
 
-        Ok(vec![ReturnSuccess::value(row)])
+            Ok(vec![ReturnSuccess::value(row)])
+        } else {
+            Ok(self
+                .values
+                .iter()
+                .map(|row| ReturnSuccess::value(row.clone()))
+                .collect::<Vec<_>>())
+        }
     }
 }
 
diff --git a/tests/filters_test.rs b/tests/filters_test.rs
index 7f38b6047..ff50a8412 100644
--- a/tests/filters_test.rs
+++ b/tests/filters_test.rs
@@ -727,8 +727,8 @@ fn can_get_reverse_first() {
 }
 
 #[test]
-fn embed() {
-    Playground::setup("embed_test", |dirs, sandbox| {
+fn embed_rows_into_a_row() {
+    Playground::setup("embed_test_1", |dirs, sandbox| {
         sandbox.with_files(vec![FileWithContentToBeTrimmed(
             "los_tres_caballeros.txt",
             r#"
@@ -756,6 +756,36 @@ fn embed() {
     })
 }
 
+#[test]
+fn embed_rows_into_a_table() {
+    Playground::setup("embed_test_2", |dirs, sandbox| {
+        sandbox.with_files(vec![FileWithContentToBeTrimmed(
+            "los_tres_caballeros.txt",
+            r#"
+                first_name,last_name
+                Andrés,Robalino
+                Jonathan,Turner
+                Yehuda,Katz
+            "#,
+        )]);
+
+        let actual = nu!(
+            cwd: dirs.test(), h::pipeline(
+            r#"
+                open los_tres_caballeros.txt
+                | from-csv
+                | get last_name
+                | embed caballero
+                | nth 2
+                | get caballero
+                | echo $it
+            "#
+        ));
+
+        assert_eq!(actual, "Katz");
+    })
+}
+
 #[test]
 fn get() {
     Playground::setup("get_test", |dirs, sandbox| {