Skip to main content

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 = match input.try_into_stream(engine_state) {
71            Ok(input) | Err(input) => input,
72        };
73        getcol(call.head, input)
74    }
75}
76
77fn getcol(head: Span, input: PipelineData) -> Result<PipelineData, ShellError> {
78    let metadata = input.metadata();
79    match input {
80        PipelineData::Empty => Ok(PipelineData::empty()),
81        PipelineData::Value(v, ..) => {
82            let span = v.span();
83            let cols = match v {
84                Value::List {
85                    vals: input_vals, ..
86                } => get_columns(&input_vals)
87                    .into_iter()
88                    .map(move |x| Value::string(x, span))
89                    .collect(),
90                Value::Custom { val, .. } => {
91                    // TODO: should we get CustomValue to expose columns in a more efficient way?
92                    // Would be nice to be able to get columns without generating the whole value
93                    let input_as_base_value = val.to_base_value(span)?;
94                    get_columns(&[input_as_base_value])
95                        .into_iter()
96                        .map(move |x| Value::string(x, span))
97                        .collect()
98                }
99                Value::Record { val, .. } => val
100                    .into_owned()
101                    .into_iter()
102                    .map(move |(x, _)| Value::string(x, head))
103                    .collect(),
104                // Propagate errors
105                Value::Error { error, .. } => return Err(*error),
106                other => {
107                    return Err(ShellError::OnlySupportsThisInputType {
108                        exp_input_type: "record or table".into(),
109                        wrong_type: other.get_type().to_string(),
110                        dst_span: head,
111                        src_span: other.span(),
112                    });
113                }
114            };
115
116            Ok(Value::list(cols, head)
117                .into_pipeline_data()
118                .set_metadata(metadata))
119        }
120        PipelineData::ListStream(stream, ..) => {
121            let values = stream.into_iter().collect::<Vec<_>>();
122            let cols = get_columns(&values)
123                .into_iter()
124                .map(|s| Value::string(s, head))
125                .collect();
126
127            Ok(Value::list(cols, head)
128                .into_pipeline_data()
129                .set_metadata(metadata))
130        }
131        PipelineData::ByteStream(stream, ..) => Err(ShellError::OnlySupportsThisInputType {
132            exp_input_type: "record or table".into(),
133            wrong_type: "byte stream".into(),
134            dst_span: head,
135            src_span: stream.span(),
136        }),
137    }
138}
139
140#[cfg(test)]
141mod test {
142    use super::*;
143
144    #[test]
145    fn test_examples() {
146        use crate::test_examples;
147
148        test_examples(Columns {})
149    }
150}