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

* from xlsx: Add parameter --sheets

* from ods: Add parameter --sheets
This commit is contained in:
Christian Menges
2021-06-10 14:44:24 +02:00
committed by GitHub
parent 9a2fe7ec0c
commit 500683831c
6 changed files with 118 additions and 10 deletions

View File

@ -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);

View File

@ -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);

View File

@ -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");
}

View File

@ -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");
}