use crate::commands::WholeStreamCommand; use crate::prelude::*; use indexmap::IndexMap; use nu_errors::ShellError; use nu_protocol::{Primitive, ReturnSuccess, Signature, SyntaxShape, UntaggedValue, Value}; use quick_xml::events::{BytesEnd, BytesStart, BytesText, Event}; use std::collections::HashSet; use std::io::Cursor; use std::io::Write; use std::iter::FromIterator; pub struct ToXML; #[derive(Deserialize)] pub struct ToXMLArgs { pretty: Option, } #[async_trait] impl WholeStreamCommand for ToXML { fn name(&self) -> &str { "to xml" } fn signature(&self) -> Signature { Signature::build("to xml").named( "pretty", SyntaxShape::Int, "Formats the XML text with the provided indentation setting", Some('p'), ) } fn usage(&self) -> &str { "Convert table into .xml text" } async fn run( &self, args: CommandArgs, registry: &CommandRegistry, ) -> Result { to_xml(args, registry).await } } pub fn add_attributes<'a>( element: &mut quick_xml::events::BytesStart<'a>, attributes: &'a IndexMap, ) { for (k, v) in attributes.iter() { element.push_attribute((k.as_str(), v.as_str())); } } pub fn get_attributes(row: &Value) -> Option> { if let UntaggedValue::Row(r) = &row.value { if let Some(v) = r.entries.get("attributes") { if let UntaggedValue::Row(a) = &v.value { let mut h = IndexMap::new(); for (k, v) in a.entries.iter() { h.insert(k.clone(), v.convert_to_string()); } return Some(h); } } } None } pub fn get_children(row: &Value) -> Option<&Vec> { if let UntaggedValue::Row(r) = &row.value { if let Some(v) = r.entries.get("children") { if let UntaggedValue::Table(t) = &v.value { return Some(t); } } } None } pub fn is_xml_row(row: &Value) -> bool { if let UntaggedValue::Row(r) = &row.value { let keys: HashSet<&String> = HashSet::from_iter(r.keys()); let children: String = "children".to_string(); let attributes: String = "attributes".to_string(); return keys.contains(&children) && keys.contains(&attributes) && keys.len() == 2; } false } pub fn write_xml_events( current: &Value, writer: &mut quick_xml::Writer, ) -> Result<(), ShellError> { match ¤t.value { UntaggedValue::Row(o) => { for (k, v) in o.entries.iter() { let mut e = BytesStart::owned(k.as_bytes(), k.len()); if !is_xml_row(v) { return Err(ShellError::labeled_error( "Expected a row with 'children' and 'attributes' columns", "missing 'children' and 'attributes' columns ", ¤t.tag, )); } let a = get_attributes(v); if let Some(ref a) = a { add_attributes(&mut e, a); } writer .write_event(Event::Start(e)) .expect("Couldn't open XML node"); let c = get_children(v); if let Some(c) = c { for v in c { write_xml_events(v, writer)?; } } writer .write_event(Event::End(BytesEnd::borrowed(k.as_bytes()))) .expect("Couldn't close XML node"); } } UntaggedValue::Table(t) => { for v in t { write_xml_events(v, writer)?; } } _ => { let s = current.convert_to_string(); writer .write_event(Event::Text(BytesText::from_plain_str(s.as_str()))) .expect("Couldn't write XML text"); } } Ok(()) } async fn to_xml(args: CommandArgs, registry: &CommandRegistry) -> Result { let registry = registry.clone(); let name_tag = args.call_info.name_tag.clone(); let name_span = name_tag.span; let (ToXMLArgs { pretty }, input) = args.process(®istry).await?; let input: Vec = input.collect().await; let to_process_input = match input.len() { x if x > 1 => { let tag = input[0].tag.clone(); vec![Value { value: UntaggedValue::Table(input), tag, }] } 1 => input, _ => vec![], }; Ok( futures::stream::iter(to_process_input.into_iter().map(move |value| { let mut w = pretty.as_ref().map_or_else( || quick_xml::Writer::new(Cursor::new(Vec::new())), |p| { quick_xml::Writer::new_with_indent( Cursor::new(Vec::new()), b' ', p.value.expect_int() as usize, ) }, ); let value_span = value.tag.span; match write_xml_events(&value, &mut w) { Ok(_) => { let b = w.into_inner().into_inner(); let s = String::from_utf8(b)?; ReturnSuccess::value( UntaggedValue::Primitive(Primitive::String(s)).into_value(&name_tag), ) } Err(_) => Err(ShellError::labeled_error_with_secondary( "Expected a table with XML-compatible structure from pipeline", "requires XML-compatible input", name_span, "originates from here".to_string(), value_span, )), } })) .to_output_stream(), ) } #[cfg(test)] mod tests { use super::ToXML; #[test] fn examples_work_as_expected() { use crate::examples::test as test_examples; test_examples(ToXML {}) } }