nushell/crates/nu-cli/src/commands/flatten.rs

use crate::command_registry::CommandRegistry;
use crate::commands::WholeStreamCommand;
use crate::prelude::*;
use nu_errors::ShellError;
use nu_protocol::{
    Dictionary, ReturnSuccess, Signature, SyntaxShape, TaggedDictBuilder, UntaggedValue, Value,
};
use nu_source::Tagged;

pub struct Command;

#[derive(Deserialize)]
pub struct Arguments {
    rest: Vec<Tagged<String>>,
}

#[async_trait]
impl WholeStreamCommand for Command {
    fn name(&self) -> &str {
        "flatten"
    }

    fn signature(&self) -> Signature {
        Signature::build("flatten").rest(SyntaxShape::String, "optionally flatten data by column")
    }

    fn usage(&self) -> &str {
        "Flatten the table."
    }

    async fn run(
        &self,
        args: CommandArgs,
        registry: &CommandRegistry,
    ) -> Result<OutputStream, ShellError> {
        flatten(args, registry).await
    }

    fn examples(&self) -> Vec<Example> {
        vec![
            Example {
                description: "flatten a table",
                example: "echo [[N, u, s, h, e, l, l]] | flatten | first",
                result: Some(vec![Value::from("N")]),
            },
            Example {
                description: "flatten a column having a nested table",
                example: "echo [[origin, people]; [Ecuador, $(echo [[name, meal]; ['Andres', 'arepa']])]] | flatten | get meal",
                result: Some(vec![Value::from("arepa")]),
            },
            Example {
                description: "restrict the flattening by passing column names",
                example: "echo [[origin, crate, versions]; [World, $(echo [[name]; ['nu-cli']]), ['0.21', '0.22']]] | flatten versions | last | = $it.versions",
                result: Some(vec![Value::from("0.22")]),
            }
        ]
    }
}

async fn flatten(
    args: CommandArgs,
    registry: &CommandRegistry,
) -> Result<OutputStream, ShellError> {
    let tag = args.call_info.name_tag.clone();
    let registry = registry.clone();
    let (Arguments { rest: columns }, input) = args.process(&registry).await?;

    Ok(input
        .map(move |item| {
            futures::stream::iter(flat_value(&columns, &item, &tag).into_iter().flatten())
        })
        .flatten()
        .to_output_stream())
}

enum TableInside<'a> {
    Entries(&'a str, &'a Tag, Vec<&'a Value>),
}

fn flat_value(
    columns: &[Tagged<String>],
    item: &Value,
    name_tag: impl Into<Tag>,
) -> Result<Vec<Result<ReturnSuccess, ShellError>>, ShellError> {
    let tag = item.tag.clone();
    let name_tag = name_tag.into();

    let res = {
        if item.is_row() {
            let mut out = TaggedDictBuilder::new(tag);
            let mut a_table = None;
            let mut tables_explicitly_flattened = 0;

            for (column, value) in item.row_entries() {
                let column_requested = columns.iter().find(|c| c.item == *column);

                if let Value {
                    value: UntaggedValue::Row(Dictionary { entries: mapa }),
                    ..
                } = value
                {
                    if column_requested.is_none() && !columns.is_empty() {
                        if out.contains_key(&column) {
                            out.insert_value(format!("{}_{}", column, column), value.clone());
                        } else {
                            out.insert_value(column, value.clone());
                        }
                        continue;
                    }

                    for (k, v) in mapa.into_iter() {
                        if out.contains_key(k) {
                            out.insert_value(format!("{}_{}", column, k), v.clone());
                        } else {
                            out.insert_value(k, v.clone());
                        }
                    }
                } else if value.is_table() {
                    if tables_explicitly_flattened >= 1 && column_requested.is_some() {
                        let attempted = if let Some(name) = column_requested {
                            name.span()
                        } else {
                            name_tag.span
                        };

                        let already_flattened =
                            if let Some(TableInside::Entries(_, column_tag, _)) = a_table {
                                column_tag.span
                            } else {
                                name_tag.span
                            };

                        return Ok(vec![ReturnSuccess::value(
                            UntaggedValue::Error(ShellError::labeled_error_with_secondary(
                                "can only flatten one inner table at the same time",
                                "tried flattening more than one column with inner tables",
                                attempted,
                                "...but is flattened already",
                                already_flattened,
                            ))
                            .into_value(name_tag),
                        )]);
                    }

                    if !columns.is_empty() {
                        if let Some(requested) = column_requested {
                            a_table = Some(TableInside::Entries(
                                &requested.item,
                                &requested.tag,
                                value.table_entries().collect(),
                            ));

                            tables_explicitly_flattened += 1;
                        } else {
                            out.insert_value(column, value.clone());
                        }
                    } else if a_table.is_none() {
                        a_table = Some(TableInside::Entries(
                            &column,
                            &value.tag,
                            value.table_entries().collect(),
                        ))
                    } else {
                        out.insert_value(column, value.clone());
                    }
                } else {
                    out.insert_value(column, value.clone());
                }
            }

            let mut expanded = vec![];

            if let Some(TableInside::Entries(column, _, entries)) = a_table {
                for entry in entries.into_iter() {
                    let mut base = out.clone();
                    base.insert_value(column, entry.clone());
                    expanded.push(base.into_value());
                }
            } else {
                expanded.push(out.into_value());
            }

            expanded
        } else if item.is_table() {
            item.table_entries().map(Clone::clone).collect()
        } else {
            vec![item.clone()]
        }
    };

    Ok(res.into_iter().map(ReturnSuccess::value).collect())
}