nu_command/matrix/
reshape.rs

1use crate::matrix::MatrixValue;
2use nu_engine::command_prelude::*;
3
4#[derive(Clone)]
5pub struct MatrixReshape;
6
7impl Command for MatrixReshape {
8    fn name(&self) -> &str {
9        "matrix reshape"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("matrix reshape")
14            .input_output_types(vec![(
15                Type::Custom("matrix".into()),
16                Type::Custom("matrix".into()),
17            )])
18            .optional(
19                "dimensions",
20                SyntaxShape::Int,
21                "The new dimensions (e.g., 2 3 for a 2x3 matrix). Required unless --flatten is used.",
22            )
23            .rest(
24                "more_dimensions",
25                SyntaxShape::Int,
26                "Additional dimensions for n-dimensional reshaping.",
27            )
28            .switch("flatten", "Flatten the matrix to a 1D vector", Some('f'))
29            .category(Category::Filters)
30    }
31
32    fn description(&self) -> &str {
33        "Change the dimensions of a matrix."
34    }
35
36    fn search_terms(&self) -> Vec<&str> {
37        vec!["dimensions", "flatten"]
38    }
39
40    fn run(
41        &self,
42        engine_state: &EngineState,
43        stack: &mut Stack,
44        call: &Call,
45        input: PipelineData,
46    ) -> Result<PipelineData, ShellError> {
47        let head = call.head;
48        let flatten = call.has_flag(engine_state, stack, "flatten")?;
49        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;
50
51        let result = if flatten {
52            let total: usize = matrix.array.len();
53            matrix
54                .array
55                .into_shape_with_order(ndarray::IxDyn(&[total]))
56                .map_err(|e| {
57                    ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
58                        "Reshape error",
59                        e.to_string(),
60                        head,
61                    ))
62                })?
63        } else {
64            let first_dim: Option<i64> = call.opt(engine_state, stack, 0)?;
65            let first_dim = first_dim.ok_or_else(|| {
66                ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
67                    "Missing dimensions",
68                    "at least one dimension is required, or use --flatten",
69                    head,
70                ))
71            })? as usize;
72            let rest_dims: Vec<Value> = call.rest(engine_state, stack, 1)?;
73
74            let mut shape = vec![first_dim];
75            for v in rest_dims {
76                match v.as_int() {
77                    Ok(d) if d > 0 => shape.push(d as usize),
78                    _ => {
79                        return Err(ShellError::Generic(
80                            nu_protocol::shell_error::generic::GenericError::new(
81                                "Invalid dimensions",
82                                "dimensions must be positive integers",
83                                head,
84                            ),
85                        ));
86                    }
87                }
88            }
89
90            for &d in &shape {
91                if d == 0 {
92                    return Err(ShellError::Generic(
93                        nu_protocol::shell_error::generic::GenericError::new(
94                            "Invalid dimensions",
95                            "dimensions must be positive",
96                            head,
97                        ),
98                    ));
99                }
100            }
101
102            let expected: usize = shape.iter().product();
103            if expected != matrix.array.len() {
104                return Err(ShellError::Generic(
105                    nu_protocol::shell_error::generic::GenericError::new(
106                        "Shape mismatch",
107                        format!(
108                            "cannot reshape {} elements into shape {:?} (would have {} elements)",
109                            matrix.array.len(),
110                            shape,
111                            expected
112                        ),
113                        head,
114                    ),
115                ));
116            }
117
118            matrix
119                .array
120                .into_shape_with_order(ndarray::IxDyn(&shape))
121                .map_err(|e| {
122                    ShellError::Generic(nu_protocol::shell_error::generic::GenericError::new(
123                        "Reshape error",
124                        e.to_string(),
125                        head,
126                    ))
127                })?
128        };
129
130        Ok(MatrixValue::new(result)
131            .into_value(head)
132            .into_pipeline_data())
133    }
134
135    fn examples(&self) -> Vec<Example<'static>> {
136        vec![
137            Example {
138                description: "Reshape a 1x6 matrix to 2x3",
139                example: "[[1 2 3 4 5 6]] | into matrix | matrix reshape 2 3 | matrix into-nu | to nuon",
140                result: Some(Value::test_string("[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]")),
141            },
142            Example {
143                description: "Flatten a 2x2 matrix to 1D",
144                example: "matrix identity 2 | matrix reshape --flatten | matrix into-nu | to nuon",
145                result: Some(Value::test_string("[[1.0, 0.0, 0.0, 1.0]]")),
146            },
147        ]
148    }
149}
150
151#[cfg(test)]
152mod test {
153    use super::*;
154
155    #[test]
156    fn test_examples() -> nu_test_support::Result {
157        nu_test_support::test().examples(MatrixReshape)
158    }
159}
nu_command/matrix/reshape.rs

nu_command/matrix/
reshape.rs
