Skip to main content

nu_command/conversions/into/
matrix.rs

1use crate::matrix::MatrixValue;
2use nu_engine::command_prelude::*;
3
4#[derive(Clone)]
5pub struct IntoMatrix;
6
7impl Command for IntoMatrix {
8    fn name(&self) -> &str {
9        "into matrix"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("into matrix")
14            .input_output_types(vec![
15                (Type::table(), Type::Custom("matrix".into())),
16                (
17                    Type::List(Box::new(Type::List(Box::new(Type::Number)))),
18                    Type::Custom("matrix".into()),
19                ),
20            ])
21            .category(Category::Conversions)
22    }
23
24    fn description(&self) -> &str {
25        "Convert a nushell table or list of lists into a matrix."
26    }
27
28    fn search_terms(&self) -> Vec<&str> {
29        vec!["convert", "array", "ndarray", "2d"]
30    }
31
32    fn run(
33        &self,
34        _engine_state: &EngineState,
35        _stack: &mut Stack,
36        call: &Call,
37        input: PipelineData,
38    ) -> Result<PipelineData, ShellError> {
39        let head = call.head;
40        let values: Vec<Value> = input.into_iter().collect();
41        into_matrix(&values, head)
42    }
43
44    fn examples(&self) -> Vec<Example<'static>> {
45        vec![
46            Example {
47                description: "Convert a list of lists to a matrix",
48                example: "[[1 2 3] [4 5 6]] | into matrix | matrix into-nu | to nuon",
49                result: Some(Value::test_string("[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]")),
50            },
51            Example {
52                description: "Convert a list of records to a matrix",
53                example: "[{a: 1 b: 2} {a: 3 b: 4}] | into matrix | matrix into-nu | to nuon",
54                result: Some(Value::test_string("[[1.0, 2.0], [3.0, 4.0]]")),
55            },
56        ]
57    }
58}
59
60fn into_matrix(values: &[Value], span: Span) -> Result<PipelineData, ShellError> {
61    if values.is_empty() {
62        let array = ndarray::ArrayD::from_shape_vec(vec![0, 0], vec![]).map_err(|e| {
63            ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
64                "Matrix shape error",
65                e.to_string(),
66                span,
67            ))
68        })?;
69        return Ok(MatrixValue::new(array)
70            .into_value(span)
71            .into_pipeline_data());
72    }
73
74    match &values[0] {
75        Value::List { .. } => {
76            let matrix = MatrixValue::from_list_of_lists(values, span)?;
77            Ok(matrix.into_value(span).into_pipeline_data())
78        }
79        Value::Record { .. } => {
80            let matrix = MatrixValue::from_list_of_records(values, span)?;
81            Ok(matrix.into_value(span).into_pipeline_data())
82        }
83        Value::Custom { val, .. } if val.type_name() == "matrix" => {
84            Ok(val.clone_value(span).into_pipeline_data())
85        }
86        _ => Err(ShellError::Generic(
87            nu_protocol::shell_error::generic::GenericError::new(
88                "Invalid input",
89                "expected a list of lists, a list of records, or a matrix",
90                span,
91            ),
92        )),
93    }
94}
95
96#[cfg(test)]
97mod test {
98    use super::*;
99
100    #[test]
101    fn test_examples() -> nu_test_support::Result {
102        nu_test_support::test().examples(IntoMatrix)
103    }
104}