from xlsx/ods: Add parameter --sheets (#3600)

* from xlsx: Add parameter --sheets

* from ods: Add parameter --sheets

crates/nu-command/src/commands/from_ods.rs

@@ -3,7 +3,7 @@ use calamine::*;
 use nu_data::TaggedListBuilder;
 use nu_engine::WholeStreamCommand;
 use nu_errors::ShellError;
-use nu_protocol::{Signature, TaggedDictBuilder, UntaggedValue};
+use nu_protocol::{Primitive, Signature, SyntaxShape, TaggedDictBuilder, UntaggedValue, Value};
 use std::io::Cursor;
 
 pub struct FromOds;
@@ -14,7 +14,12 @@ impl WholeStreamCommand for FromOds {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build("from ods")
+        Signature::build("from ods").named(
+            "sheets",
+            SyntaxShape::Table,
+            "Only convert specified sheets",
+            Some('s'),
+        )
     }
 
     fn usage(&self) -> &str {
@@ -26,10 +31,37 @@ impl WholeStreamCommand for FromOds {
     }
 }
 
+// Adapted from crates/nu-command/src/commands/dataframe/utils.rs
+fn convert_columns(columns: &[Value]) -> Result<Vec<String>, ShellError> {
+    let res = columns
+        .iter()
+        .map(|value| match &value.value {
+            UntaggedValue::Primitive(Primitive::String(s)) => Ok(s.clone()),
+            _ => Err(ShellError::labeled_error(
+                "Incorrect column format",
+                "Only string as column name",
+                &value.tag,
+            )),
+        })
+        .collect::<Result<Vec<String>, _>>()?;
+
+    Ok(res)
+}
+
 fn from_ods(args: CommandArgs) -> Result<OutputStream, ShellError> {
     let tag = args.call_info.name_tag.clone();
     let span = tag.span;
 
+    let args = args.evaluate_once()?;
+
+    let mut sel_sheets = vec![];
+
+    if let Some(s) = args.call_info.args.get("sheets") {
+        if let UntaggedValue::Table(columns) = s.value.clone() {
+            sel_sheets = convert_columns(columns.as_slice())?;
+        }
+    }
+
     let bytes = args.input.collect_binary(tag.clone())?;
     let buf: Cursor<Vec<u8>> = Cursor::new(bytes.item);
     let mut ods = Ods::<_>::new(buf).map_err(|_| {
@@ -38,7 +70,10 @@ fn from_ods(args: CommandArgs) -> Result<OutputStream, ShellError> {
 
     let mut dict = TaggedDictBuilder::new(&tag);
 
-    let sheet_names = ods.sheet_names().to_owned();
+    let mut sheet_names = ods.sheet_names().to_owned();
+    if !sel_sheets.is_empty() {
+        sheet_names.retain(|e| sel_sheets.contains(e));
+    }
 
     for sheet_name in &sheet_names {
         let mut sheet_output = TaggedListBuilder::new(&tag);

crates/nu-command/src/commands/from_xlsx.rs

@@ -3,7 +3,7 @@ use calamine::*;
 use nu_data::TaggedListBuilder;
 use nu_engine::WholeStreamCommand;
 use nu_errors::ShellError;
-use nu_protocol::{Signature, TaggedDictBuilder, UntaggedValue};
+use nu_protocol::{Primitive, Signature, SyntaxShape, TaggedDictBuilder, UntaggedValue, Value};
 use std::io::Cursor;
 
 pub struct FromXlsx;
@@ -14,11 +14,18 @@ impl WholeStreamCommand for FromXlsx {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build("from xlsx").switch(
-            "noheaders",
-            "don't treat the first row as column names",
-            Some('n'),
-        )
+        Signature::build("from xlsx")
+            .switch(
+                "noheaders",
+                "don't treat the first row as column names",
+                Some('n'),
+            )
+            .named(
+                "sheets",
+                SyntaxShape::Table,
+                "Only convert specified sheets",
+                Some('s'),
+            )
     }
 
     fn usage(&self) -> &str {
@@ -30,10 +37,37 @@ impl WholeStreamCommand for FromXlsx {
     }
 }
 
+// Adapted from crates/nu-command/src/commands/dataframe/utils.rs
+fn convert_columns(columns: &[Value]) -> Result<Vec<String>, ShellError> {
+    let res = columns
+        .iter()
+        .map(|value| match &value.value {
+            UntaggedValue::Primitive(Primitive::String(s)) => Ok(s.clone()),
+            _ => Err(ShellError::labeled_error(
+                "Incorrect column format",
+                "Only string as column name",
+                &value.tag,
+            )),
+        })
+        .collect::<Result<Vec<String>, _>>()?;
+
+    Ok(res)
+}
+
 fn from_xlsx(args: CommandArgs) -> Result<OutputStream, ShellError> {
     let tag = args.call_info.name_tag.clone();
     let span = tag.span;
 
+    let args = args.evaluate_once()?;
+
+    let mut sel_sheets = vec![];
+
+    if let Some(s) = args.call_info.args.get("sheets") {
+        if let UntaggedValue::Table(columns) = s.value.clone() {
+            sel_sheets = convert_columns(columns.as_slice())?;
+        }
+    }
+
     let value = args.input.collect_binary(tag.clone())?;
 
     let buf: Cursor<Vec<u8>> = Cursor::new(value.item);
@@ -43,7 +77,10 @@ fn from_xlsx(args: CommandArgs) -> Result<OutputStream, ShellError> {
 
     let mut dict = TaggedDictBuilder::new(&tag);
 
-    let sheet_names = xls.sheet_names().to_owned();
+    let mut sheet_names = xls.sheet_names().to_owned();
+    if !sel_sheets.is_empty() {
+        sheet_names.retain(|e| sel_sheets.contains(e));
+    }
 
     for sheet_name in &sheet_names {
         let mut sheet_output = TaggedListBuilder::new(&tag);

crates/nu-command/tests/format_conversions/ods.rs

@@ -14,3 +14,17 @@ fn from_ods_file_to_table() {
 
     assert_eq!(actual.out, "Gill");
 }
+
+#[test]
+fn from_ods_file_to_table_select_sheet() {
+    let actual = nu!(
+        cwd: "tests/fixtures/formats", pipeline(
+        r#"
+            open sample_data.ods --raw
+            | from ods -s ["SalesOrders"]
+            | get
+        "#
+    ));
+
+    assert_eq!(actual.out, "SalesOrders");
+}

crates/nu-command/tests/format_conversions/xlsx.rs

@@ -14,3 +14,17 @@ fn from_excel_file_to_table() {
 
     assert_eq!(actual.out, "Gill");
 }
+
+#[test]
+fn from_excel_file_to_table_select_sheet() {
+    let actual = nu!(
+        cwd: "tests/fixtures/formats", pipeline(
+        r#"
+            open sample_data.xlsx --raw
+            | from xlsx -s ["SalesOrders"]
+            | get
+        "#
+    ));
+
+    assert_eq!(actual.out, "SalesOrders");
+}