shannon_nu_command/filters/
columns.rs

1use nu_engine::{column::get_columns, command_prelude::*};
2
3#[derive(Clone)]
4pub struct Columns;
5
6impl Command for Columns {
7    fn name(&self) -> &str {
8        "columns"
9    }
10
11    fn signature(&self) -> Signature {
12        Signature::build(self.name())
13            .input_output_types(vec![
14                (Type::table(), Type::List(Box::new(Type::String))),
15                (Type::record(), Type::List(Box::new(Type::String))),
16            ])
17            .category(Category::Filters)
18    }
19
20    fn description(&self) -> &str {
21        "Given a record or table, produce a list of its columns' names."
22    }
23
24    fn extra_description(&self) -> &str {
25        "This is a counterpart to `values`, which produces a list of columns' values."
26    }
27
28    fn examples(&self) -> Vec<Example<'_>> {
29        vec![
30            Example {
31                example: "{ acronym:PWD, meaning:'Print Working Directory' } | columns",
32                description: "Get the columns from the record",
33                result: Some(Value::list(
34                    vec![Value::test_string("acronym"), Value::test_string("meaning")],
35                    Span::test_data(),
36                )),
37            },
38            Example {
39                example: "[[name,age,grade]; [bill,20,a]] | columns",
40                description: "Get the columns from the table",
41                result: Some(Value::list(
42                    vec![
43                        Value::test_string("name"),
44                        Value::test_string("age"),
45                        Value::test_string("grade"),
46                    ],
47                    Span::test_data(),
48                )),
49            },
50            Example {
51                example: "[[name,age,grade]; [bill,20,a]] | columns | first",
52                description: "Get the first column from the table",
53                result: None,
54            },
55            Example {
56                example: "[[name,age,grade]; [bill,20,a]] | columns | select 1",
57                description: "Get the second column from the table",
58                result: None,
59            },
60        ]
61    }
62
63    fn run(
64        &self,
65        engine_state: &EngineState,
66        _stack: &mut Stack,
67        call: &Call,
68        input: PipelineData,
69    ) -> Result<PipelineData, ShellError> {
70        let input = input.into_stream_or_original(engine_state);
71        getcol(call.head, input)
72    }
73}
74
75fn getcol(head: Span, input: PipelineData) -> Result<PipelineData, ShellError> {
76    let metadata = input.metadata();
77    match input {
78        PipelineData::Empty => Ok(PipelineData::empty()),
79        PipelineData::Value(v, ..) => {
80            let span = v.span();
81            let cols = match v {
82                Value::List {
83                    vals: input_vals, ..
84                } => get_columns(&input_vals)
85                    .into_iter()
86                    .map(move |x| Value::string(x, span))
87                    .collect(),
88                Value::Custom { val, .. } => {
89                    // TODO: should we get CustomValue to expose columns in a more efficient way?
90                    // Would be nice to be able to get columns without generating the whole value
91                    let input_as_base_value = val.to_base_value(span)?;
92                    get_columns(&[input_as_base_value])
93                        .into_iter()
94                        .map(move |x| Value::string(x, span))
95                        .collect()
96                }
97                Value::Record { val, .. } => val
98                    .into_owned()
99                    .into_iter()
100                    .map(move |(x, _)| Value::string(x, head))
101                    .collect(),
102                // Propagate errors
103                Value::Error { error, .. } => return Err(*error),
104                other => {
105                    return Err(ShellError::OnlySupportsThisInputType {
106                        exp_input_type: "record or table".into(),
107                        wrong_type: other.get_type().to_string(),
108                        dst_span: head,
109                        src_span: other.span(),
110                    });
111                }
112            };
113
114            Ok(Value::list(cols, head)
115                .into_pipeline_data()
116                .set_metadata(metadata))
117        }
118        PipelineData::ListStream(stream, ..) => {
119            let values = stream.into_iter().collect::<Vec<_>>();
120            let cols = get_columns(&values)
121                .into_iter()
122                .map(|s| Value::string(s, head))
123                .collect();
124
125            Ok(Value::list(cols, head)
126                .into_pipeline_data()
127                .set_metadata(metadata))
128        }
129        PipelineData::ByteStream(stream, ..) => Err(ShellError::OnlySupportsThisInputType {
130            exp_input_type: "record or table".into(),
131            wrong_type: "byte stream".into(),
132            dst_span: head,
133            src_span: stream.span(),
134        }),
135    }
136}
137
138#[cfg(test)]
139mod test {
140    use super::*;
141
142    #[test]
143    fn test_examples() -> nu_test_support::Result {
144        nu_test_support::test().examples(Columns)
145    }
146}
shannon_nu_command/filters/columns.rs

shannon_nu_command/filters/
columns.rs
