diff --git a/crates/nu-command/src/commands/from_ods.rs b/crates/nu-command/src/commands/from_ods.rs index c39fb29852..f8cf3c705a 100644 --- a/crates/nu-command/src/commands/from_ods.rs +++ b/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, 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::, _>>()?; + + Ok(res) +} + fn from_ods(args: CommandArgs) -> Result { 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> = Cursor::new(bytes.item); let mut ods = Ods::<_>::new(buf).map_err(|_| { @@ -38,7 +70,10 @@ fn from_ods(args: CommandArgs) -> Result { 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); diff --git a/crates/nu-command/src/commands/from_xlsx.rs b/crates/nu-command/src/commands/from_xlsx.rs index a6ba395832..df037f8900 100644 --- a/crates/nu-command/src/commands/from_xlsx.rs +++ b/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, 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::, _>>()?; + + Ok(res) +} + fn from_xlsx(args: CommandArgs) -> Result { 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> = Cursor::new(value.item); @@ -43,7 +77,10 @@ fn from_xlsx(args: CommandArgs) -> Result { 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); diff --git a/crates/nu-command/tests/format_conversions/ods.rs b/crates/nu-command/tests/format_conversions/ods.rs index dce6042713..5d5ce7ca7a 100644 --- a/crates/nu-command/tests/format_conversions/ods.rs +++ b/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"); +} diff --git a/crates/nu-command/tests/format_conversions/xlsx.rs b/crates/nu-command/tests/format_conversions/xlsx.rs index 5fffb8376f..36b3aca17a 100644 --- a/crates/nu-command/tests/format_conversions/xlsx.rs +++ b/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"); +} diff --git a/docs/commands/from-ods.md b/docs/commands/from-ods.md index f04f9aa324..ef7e808fc3 100644 --- a/docs/commands/from-ods.md +++ b/docs/commands/from-ods.md @@ -2,6 +2,10 @@ Parses OpenDocument Spreadsheet binary data into a table. `open` calls `from ods` automatically when the file extension is `ods`. Use this command when `open` is unable to guess the file type from the extension. +## Flags +* -h, --help: Display this help message +* -s, --sheets \[\ \ ... \]: Only convert specified sheets. Non-existing sheets are skipped. + ## Examples ```sh diff --git a/docs/commands/from-xlsx.md b/docs/commands/from-xlsx.md index 9245225763..0f52d4cb67 100644 --- a/docs/commands/from-xlsx.md +++ b/docs/commands/from-xlsx.md @@ -2,6 +2,10 @@ Parses MS Excel binary data into a table. `open` calls `from xlsx` automatically when the file extension is `xlsx`. Use this command when `open` is unable to guess the file type from the extension. +## Flags +* -h, --help: Display this help message +* -s, --sheets \[\ \ ... \]: Only convert specified sheets. Non-existing sheets are skipped. + ## Examples ```shell