nu_command/matrix/
get_col.rs

1use crate::matrix::MatrixValue;
2use nu_engine::command_prelude::*;
3
4#[derive(Clone)]
5pub struct MatrixGetCol;
6
7impl Command for MatrixGetCol {
8    fn name(&self) -> &str {
9        "matrix get-col"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("matrix get-col")
14            .input_output_types(vec![(
15                Type::Custom("matrix".into()),
16                Type::List(Box::new(Type::Float)),
17            )])
18            .required(
19                "index",
20                SyntaxShape::Int,
21                "The column index to extract (0-based).",
22            )
23            .category(Category::Filters)
24    }
25
26    fn description(&self) -> &str {
27        "Extract a column from a 2D matrix."
28    }
29
30    fn search_terms(&self) -> Vec<&str> {
31        vec!["column"]
32    }
33
34    fn run(
35        &self,
36        engine_state: &EngineState,
37        stack: &mut Stack,
38        call: &Call,
39        input: PipelineData,
40    ) -> Result<PipelineData, ShellError> {
41        let head = call.head;
42        let index: usize = call.req::<i64>(engine_state, stack, 0)? as usize;
43        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;
44
45        if matrix.array.ndim() != 2 {
46            return Err(ShellError::Generic(
47                nu_protocol::shell_error::generic::GenericError::new(
48                    "Invalid dimensions",
49                    format!(
50                        "get-col requires a 2D matrix, got {} dimensions",
51                        matrix.array.ndim()
52                    ),
53                    head,
54                ),
55            ));
56        }
57
58        let ncols = matrix.array.shape()[1];
59        if index >= ncols {
60            return Err(ShellError::Generic(
61                nu_protocol::shell_error::generic::GenericError::new(
62                    "Index out of bounds",
63                    format!(
64                        "column index {} is out of bounds, matrix has {} columns",
65                        index, ncols
66                    ),
67                    head,
68                ),
69            ));
70        }
71
72        let col = matrix.array.index_axis(ndarray::Axis(1), index);
73        let vals: Vec<Value> = col.iter().map(|v| Value::float(*v, head)).collect();
74        Ok(Value::list(vals, head).into_pipeline_data())
75    }
76
77    fn examples(&self) -> Vec<Example<'static>> {
78        vec![
79            Example {
80                description: "Get the first column of a 2x3 matrix",
81                example: "matrix zeros 2 3 | matrix set-col 0 [1.0 2.0] | matrix get-col 0 | to nuon",
82                result: Some(Value::test_string("[1.0, 2.0]")),
83            },
84            Example {
85                description: "Get the second column of a 2x2 identity matrix",
86                example: "matrix identity 2 | matrix get-col 1 | to nuon",
87                result: Some(Value::test_string("[0.0, 1.0]")),
88            },
89        ]
90    }
91}
92
93#[cfg(test)]
94mod test {
95    use super::*;
96
97    #[test]
98    fn test_examples() -> nu_test_support::Result {
99        nu_test_support::test().examples(MatrixGetCol)
100    }
101}
nu_command/matrix/get_col.rs

nu_command/matrix/
get_col.rs
