nu_command/matrix/
set_row.rs

1use crate::matrix::MatrixValue;
2use crate::matrix::value::values_to_f64s;
3use nu_engine::command_prelude::*;
4
5#[derive(Clone)]
6pub struct MatrixSetRow;
7
8impl Command for MatrixSetRow {
9    fn name(&self) -> &str {
10        "matrix set-row"
11    }
12
13    fn signature(&self) -> Signature {
14        Signature::build("matrix set-row")
15            .input_output_types(vec![(
16                Type::Custom("matrix".into()),
17                Type::Custom("matrix".into()),
18            )])
19            .required(
20                "index",
21                SyntaxShape::Int,
22                "The row index to replace (0-based).",
23            )
24            .required(
25                "replacement",
26                SyntaxShape::List(Box::new(SyntaxShape::Number)),
27                "The new row values as a list of numbers.",
28            )
29            .category(Category::Filters)
30    }
31
32    fn description(&self) -> &str {
33        "Replace a row in a matrix."
34    }
35
36    fn search_terms(&self) -> Vec<&str> {
37        vec!["replace"]
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 index: usize = call.req::<i64>(engine_state, stack, 0)? as usize;
49        let replacement: Value = call.req(engine_state, stack, 1)?;
50        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;
51
52        let replacement_vals = match replacement {
53            Value::List { vals, .. } => vals,
54            _ => {
55                return Err(ShellError::Generic(
56                    nu_protocol::shell_error::generic::GenericError::new(
57                        "Invalid replacement",
58                        "expected a list of numbers",
59                        head,
60                    ),
61                ));
62            }
63        };
64
65        let ncols = if matrix.array.ndim() >= 2 {
66            matrix.array.shape()[1]
67        } else {
68            1
69        };
70        if replacement_vals.len() != ncols {
71            return Err(ShellError::Generic(
72                nu_protocol::shell_error::generic::GenericError::new(
73                    "Size mismatch",
74                    format!(
75                        "replacement has {} elements, but row has {} columns",
76                        replacement_vals.len(),
77                        ncols
78                    ),
79                    head,
80                ),
81            ));
82        }
83
84        if index >= matrix.array.shape()[0] {
85            return Err(ShellError::Generic(
86                nu_protocol::shell_error::generic::GenericError::new(
87                    "Index out of bounds",
88                    format!(
89                        "row index {} is out of bounds, matrix has {} rows",
90                        index,
91                        matrix.array.shape()[0]
92                    ),
93                    head,
94                ),
95            ));
96        }
97
98        let floats = values_to_f64s(&replacement_vals, head)?;
99
100        let mut new_array = matrix.array;
101        for (j, &val) in floats.iter().enumerate() {
102            if new_array.ndim() == 1 {
103                new_array[[index]] = val;
104            } else {
105                new_array[[index, j]] = val;
106            }
107        }
108
109        Ok(MatrixValue::new(new_array)
110            .into_value(head)
111            .into_pipeline_data())
112    }
113
114    fn examples(&self) -> Vec<Example<'static>> {
115        vec![
116            Example {
117                description: "Replace the first row of a 2x3 matrix",
118                example: "matrix zeros 2 3 | matrix set-row 0 [1.0 2.0 3.0] | matrix into-nu | to nuon",
119                result: Some(Value::test_string("[[1.0, 2.0, 3.0], [0.0, 0.0, 0.0]]")),
120            },
121            Example {
122                description: "Replace the second row of a 2x2 matrix",
123                example: "matrix zeros 2 2 | matrix set-row 1 [4.0 5.0] | matrix into-nu | to nuon",
124                result: Some(Value::test_string("[[0.0, 0.0], [4.0, 5.0]]")),
125            },
126        ]
127    }
128}
129
130#[cfg(test)]
131mod test {
132    use super::*;
133
134    #[test]
135    fn test_examples() -> nu_test_support::Result {
136        nu_test_support::test().examples(MatrixSetRow)
137    }
138}
nu_command/matrix/set_row.rs

nu_command/matrix/
set_row.rs
