nu_command/matrix/
into_nu.rs

1use crate::matrix::MatrixValue;
2use nu_engine::command_prelude::*;
3
4#[derive(Clone)]
5pub struct MatrixIntoNu;
6
7impl Command for MatrixIntoNu {
8    fn name(&self) -> &str {
9        "matrix into-nu"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("matrix into-nu")
14            .input_output_types(vec![(Type::Custom("matrix".into()), Type::table())])
15            .switch(
16                "as-records",
17                "Output as a list of records with auto-generated column names",
18                Some('r'),
19            )
20            .category(Category::Conversions)
21    }
22
23    fn description(&self) -> &str {
24        "Convert a matrix to a nushell table (list of lists by default)."
25    }
26
27    fn search_terms(&self) -> Vec<&str> {
28        vec!["convert", "table", "list"]
29    }
30
31    fn run(
32        &self,
33        engine_state: &EngineState,
34        stack: &mut Stack,
35        call: &Call,
36        input: PipelineData,
37    ) -> Result<PipelineData, ShellError> {
38        let head = call.head;
39        let as_records = call.has_flag(engine_state, stack, "as-records")?;
40        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;
41
42        if matrix.array.ndim() == 0 {
43            return Ok(
44                Value::float(matrix.array.first().copied().unwrap_or(0.0), head)
45                    .into_pipeline_data(),
46            );
47        }
48
49        if as_records {
50            matrix_to_records(&matrix.array, head)
51        } else {
52            matrix_to_lists(&matrix.array, head)
53        }
54    }
55
56    fn examples(&self) -> Vec<Example<'static>> {
57        vec![
58            Example {
59                description: "Convert an identity matrix to a nushell table",
60                example: "matrix identity 2 | matrix into-nu | to nuon",
61                result: Some(Value::test_string("[[1.0, 0.0], [0.0, 1.0]]")),
62            },
63            Example {
64                description: "Convert a matrix to records",
65                example: "matrix identity 2 | matrix into-nu --as-records | to nuon",
66                result: Some(Value::test_string(
67                    "[[\"col0\", \"col1\"]; [1.0, 0.0], [0.0, 1.0]]",
68                )),
69            },
70        ]
71    }
72}
73
74fn matrix_to_lists(array: &ndarray::ArrayD<f64>, span: Span) -> Result<PipelineData, ShellError> {
75    let rows = match array.ndim() {
76        0 => {
77            vec![Value::float(array.first().copied().unwrap_or(0.0), span)]
78        }
79        1 => {
80            let vals: Vec<Value> = array.iter().map(|v| Value::float(*v, span)).collect();
81            vec![Value::list(vals, span)]
82        }
83        2 => array
84            .axis_iter(ndarray::Axis(0))
85            .map(|row| {
86                let vals: Vec<Value> = row.iter().map(|v| Value::float(*v, span)).collect();
87                Value::list(vals, span)
88            })
89            .collect(),
90        _ => array
91            .axis_iter(ndarray::Axis(0))
92            .map(|sub| {
93                let sub_lists: Vec<Value> = sub
94                    .axis_iter(ndarray::Axis(0))
95                    .map(|inner| {
96                        let vals: Vec<Value> =
97                            inner.iter().map(|v| Value::float(*v, span)).collect();
98                        Value::list(vals, span)
99                    })
100                    .collect();
101                Value::list(sub_lists, span)
102            })
103            .collect(),
104    };
105
106    Ok(Value::list(rows, span).into_pipeline_data())
107}
108
109fn matrix_to_records(array: &ndarray::ArrayD<f64>, span: Span) -> Result<PipelineData, ShellError> {
110    let ncols = if array.ndim() >= 2 {
111        array.shape()[array.ndim() - 1]
112    } else {
113        array.len()
114    };
115    let col_names: Vec<String> = (0..ncols).map(|i| format!("col{}", i)).collect();
116
117    let rows: Vec<Value> = array
118        .axis_iter(ndarray::Axis(0))
119        .map(|row| {
120            let mut record = nu_protocol::Record::new();
121            for (j, name) in col_names.iter().enumerate() {
122                let val = if j < row.len() {
123                    if let Some(&v) = row.get(j) {
124                        Value::float(v, span)
125                    } else {
126                        Value::float(0.0, span)
127                    }
128                } else {
129                    Value::float(0.0, span)
130                };
131                record.push(name, val);
132            }
133            Value::record(record, span)
134        })
135        .collect();
136
137    Ok(Value::list(rows, span).into_pipeline_data())
138}
nu_command/matrix/into_nu.rs

nu_command/matrix/
into_nu.rs
