nushell/src/commands/from_csv.rs

122 lines
3.4 KiB
Rust
Raw Normal View History

use crate::commands::WholeStreamCommand;
2019-08-01 03:58:42 +02:00
use crate::object::{Primitive, TaggedDictBuilder, Value};
use crate::prelude::*;
use csv::ReaderBuilder;
pub struct FromCSV;
impl WholeStreamCommand for FromCSV {
fn run(
&self,
args: CommandArgs,
registry: &CommandRegistry,
) -> Result<OutputStream, ShellError> {
from_csv(args, registry)
}
fn name(&self) -> &str {
"from-csv"
}
fn signature(&self) -> Signature {
Signature::build("from-csv")
}
}
pub fn from_csv_string_to_value(
s: String,
tag: impl Into<Tag>,
2019-08-21 08:39:57 +02:00
) -> Result<Tagged<Value>, csv::Error> {
let mut reader = ReaderBuilder::new()
.has_headers(false)
.from_reader(s.as_bytes());
let tag = tag.into();
let mut fields: VecDeque<String> = VecDeque::new();
let mut iter = reader.records();
2019-07-20 08:44:21 +02:00
let mut rows = vec![];
if let Some(result) = iter.next() {
let line = result?;
for item in line.iter() {
fields.push_back(item.to_string());
}
}
loop {
if let Some(row_values) = iter.next() {
let row_values = row_values?;
let mut row = TaggedDictBuilder::new(tag);
for (idx, entry) in row_values.iter().enumerate() {
2019-08-01 03:58:42 +02:00
row.insert_tagged(
fields.get(idx).unwrap(),
Value::Primitive(Primitive::String(String::from(entry))).tagged(tag),
);
}
2019-08-01 03:58:42 +02:00
rows.push(row.into_tagged_value());
} else {
break;
}
}
Ok(Tagged::from_item(Value::List(rows), tag))
}
fn from_csv(args: CommandArgs, registry: &CommandRegistry) -> Result<OutputStream, ShellError> {
2019-07-24 00:22:11 +02:00
let args = args.evaluate_once(registry)?;
let span = args.name_span();
2019-08-21 08:39:57 +02:00
let input = args.input;
2019-08-21 08:39:57 +02:00
let stream = async_stream_block! {
let values: Vec<Tagged<Value>> = input.values.collect().await;
let mut concat_string = String::new();
let mut latest_tag: Option<Tag> = None;
for value in values {
let value_tag = value.tag();
latest_tag = Some(value_tag);
match value.item {
2019-08-01 03:58:42 +02:00
Value::Primitive(Primitive::String(s)) => {
2019-08-21 08:39:57 +02:00
concat_string.push_str(&s);
concat_string.push_str("\n");
2019-08-01 03:58:42 +02:00
}
2019-08-21 08:39:57 +02:00
_ => yield Err(ShellError::labeled_error_with_secondary(
"Expected a string from pipeline",
"requires string input",
span,
"value originates from here",
2019-08-21 08:39:57 +02:00
value_tag.span,
)),
2019-08-21 08:39:57 +02:00
2019-08-01 03:58:42 +02:00
}
2019-08-21 08:39:57 +02:00
}
match from_csv_string_to_value(concat_string, span) {
2019-08-24 09:38:38 +02:00
Ok(x) => match x {
Tagged { item: Value::List(list), .. } => {
for l in list {
yield ReturnSuccess::value(l);
}
}
x => yield ReturnSuccess::value(x),
},
2019-08-21 08:39:57 +02:00
Err(_) => if let Some(last_tag) = latest_tag {
yield Err(ShellError::labeled_error_with_secondary(
"Could not parse as CSV",
"input cannot be parsed as CSV",
span,
"value originates from here",
last_tag.span,
))
} ,
}
};
Ok(stream.to_output_stream())
}