nu-command 0.115.1

Nushell's built-in commands
Documentation
use crate::matrix::MatrixValue;
use nu_engine::command_prelude::*;

#[derive(Clone)]
pub struct MatrixIntoNu;

impl Command for MatrixIntoNu {
    fn name(&self) -> &str {
        "matrix into-nu"
    }

    fn signature(&self) -> Signature {
        Signature::build("matrix into-nu")
            .input_output_types(vec![(Type::Custom("matrix".into()), Type::table())])
            .switch(
                "as-records",
                "Output as a list of records with auto-generated column names",
                Some('r'),
            )
            .category(Category::Conversions)
    }

    fn description(&self) -> &str {
        "Convert a matrix to a nushell table (list of lists by default)."
    }

    fn search_terms(&self) -> Vec<&str> {
        vec!["convert", "table", "list"]
    }

    fn run(
        &self,
        engine_state: &EngineState,
        stack: &mut Stack,
        call: &Call,
        input: PipelineData,
    ) -> Result<PipelineData, ShellError> {
        let head = call.head;
        let as_records = call.has_flag(engine_state, stack, "as-records")?;
        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;

        if matrix.array.ndim() == 0 {
            return Ok(
                Value::float(matrix.array.first().copied().unwrap_or(0.0), head)
                    .into_pipeline_data(),
            );
        }

        if as_records {
            matrix_to_records(&matrix.array, head)
        } else {
            matrix_to_lists(&matrix.array, head)
        }
    }

    fn examples(&self) -> Vec<Example<'static>> {
        vec![
            Example {
                description: "Convert an identity matrix to a nushell table",
                example: "matrix identity 2 | matrix into-nu | to nuon",
                result: Some(Value::test_string("[[1.0, 0.0], [0.0, 1.0]]")),
            },
            Example {
                description: "Convert a matrix to records",
                example: "matrix identity 2 | matrix into-nu --as-records | to nuon",
                result: Some(Value::test_string(
                    "[[\"col0\", \"col1\"]; [1.0, 0.0], [0.0, 1.0]]",
                )),
            },
        ]
    }
}

fn matrix_to_lists(array: &ndarray::ArrayD<f64>, span: Span) -> Result<PipelineData, ShellError> {
    let rows = match array.ndim() {
        0 => {
            vec![Value::float(array.first().copied().unwrap_or(0.0), span)]
        }
        1 => {
            let vals: Vec<Value> = array.iter().map(|v| Value::float(*v, span)).collect();
            vec![Value::list(vals, span)]
        }
        2 => array
            .axis_iter(ndarray::Axis(0))
            .map(|row| {
                let vals: Vec<Value> = row.iter().map(|v| Value::float(*v, span)).collect();
                Value::list(vals, span)
            })
            .collect(),
        _ => array
            .axis_iter(ndarray::Axis(0))
            .map(|sub| {
                let sub_lists: Vec<Value> = sub
                    .axis_iter(ndarray::Axis(0))
                    .map(|inner| {
                        let vals: Vec<Value> =
                            inner.iter().map(|v| Value::float(*v, span)).collect();
                        Value::list(vals, span)
                    })
                    .collect();
                Value::list(sub_lists, span)
            })
            .collect(),
    };

    Ok(Value::list(rows, span).into_pipeline_data())
}

fn matrix_to_records(array: &ndarray::ArrayD<f64>, span: Span) -> Result<PipelineData, ShellError> {
    let ncols = if array.ndim() >= 2 {
        array.shape()[array.ndim() - 1]
    } else {
        array.len()
    };
    let col_names: Vec<String> = (0..ncols).map(|i| format!("col{}", i)).collect();

    let rows: Vec<Value> = array
        .axis_iter(ndarray::Axis(0))
        .map(|row| {
            let mut record = nu_protocol::Record::new();
            for (j, name) in col_names.iter().enumerate() {
                let val = if j < row.len() {
                    if let Some(&v) = row.get(j) {
                        Value::float(v, span)
                    } else {
                        Value::float(0.0, span)
                    }
                } else {
                    Value::float(0.0, span)
                };
                record.push(name, val);
            }
            Value::record(record, span)
        })
        .collect();

    Ok(Value::list(rows, span).into_pipeline_data())
}