headers command (#4414)

* headers command

* correct behaviour headers

crates/nu-command/src/default_context.rs

@@ -76,6 +76,7 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
             Flatten,
             Get,
             GroupBy,
+            Headers,
             SplitBy,
             Keep,
             Merge,

crates/nu-command/src/example_test.rs

@@ -13,8 +13,8 @@ use crate::To;
 
 #[cfg(test)]
 use super::{
-    Ansi, Date, From, If, Into, Math, Path, Random, Split, Str, StrCollect, StrFindReplace,
+    Ansi, Date, From, If, Into, Math, Path, Random, Split, SplitColumn, SplitRow, Str, StrCollect,
-    StrLength, Url, Wrap,
+    StrFindReplace, StrLength, Url, Wrap,
 };
 
 #[cfg(test)]
@@ -39,6 +39,8 @@ pub fn test_examples(cmd: impl Command + 'static) {
         working_set.add_decl(Box::new(Into));
         working_set.add_decl(Box::new(Random));
         working_set.add_decl(Box::new(Split));
+        working_set.add_decl(Box::new(SplitColumn));
+        working_set.add_decl(Box::new(SplitRow));
         working_set.add_decl(Box::new(Math));
         working_set.add_decl(Box::new(Path));
         working_set.add_decl(Box::new(Date));

crates/nu-command/src/filters/headers.rs

@@ -0,0 +1,148 @@
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{
+    Category, Config, Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, Value,
+};
+
+#[derive(Clone)]
+pub struct Headers;
+
+impl Command for Headers {
+    fn name(&self) -> &str {
+        "headers"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name()).category(Category::Filters)
+    }
+
+    fn usage(&self) -> &str {
+        "Use the first row of the table as column names."
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        let columns = vec!["a".to_string(), "b".to_string(), "c".to_string()];
+        vec![
+            Example {
+                description: "Returns headers from table",
+                example: r#""a b c|1 2 3" | split row "|" | split column " " | headers"#,
+                result: Some(Value::List {
+                    vals: vec![Value::Record {
+                        cols: columns.clone(),
+                        vals: vec![
+                            Value::test_string("1"),
+                            Value::test_string("2"),
+                            Value::test_string("3"),
+                        ],
+                        span: Span::test_data(),
+                    }],
+                    span: Span::test_data(),
+                }),
+            },
+            Example {
+                description: "Don't panic on rows with different headers",
+                example: r#""a b c|1 2 3|1 2 3 4" | split row "|" | split column " " | headers"#,
+                result: Some(Value::List {
+                    vals: vec![
+                        Value::Record {
+                            cols: columns.clone(),
+                            vals: vec![
+                                Value::test_string("1"),
+                                Value::test_string("2"),
+                                Value::test_string("3"),
+                            ],
+                            span: Span::test_data(),
+                        },
+                        Value::Record {
+                            cols: columns,
+                            vals: vec![
+                                Value::test_string("1"),
+                                Value::test_string("2"),
+                                Value::test_string("3"),
+                            ],
+                            span: Span::test_data(),
+                        },
+                    ],
+                    span: Span::test_data(),
+                }),
+            },
+        ]
+    }
+
+    fn run(
+        &self,
+        _engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, ShellError> {
+        let config = stack.get_config()?;
+        let value = input.into_value(call.head);
+        let headers = extract_headers(&value, &config)?;
+        let new_headers = replace_headers(value, &headers)?;
+
+        Ok(new_headers.into_pipeline_data())
+    }
+}
+
+fn replace_headers(value: Value, headers: &[String]) -> Result<Value, ShellError> {
+    match value {
+        Value::Record { vals, span, .. } => {
+            let vals = vals.into_iter().take(headers.len()).collect();
+            Ok(Value::Record {
+                cols: headers.to_owned(),
+                vals,
+                span,
+            })
+        }
+        Value::List { vals, span } => {
+            let vals = vals
+                .into_iter()
+                .skip(1)
+                .map(|value| replace_headers(value, headers))
+                .collect::<Result<Vec<Value>, ShellError>>()?;
+
+            Ok(Value::List { vals, span })
+        }
+        _ => Err(ShellError::TypeMismatch(
+            "record".to_string(),
+            value.span()?,
+        )),
+    }
+}
+
+fn extract_headers(value: &Value, config: &Config) -> Result<Vec<String>, ShellError> {
+    match value {
+        Value::Record { vals, .. } => Ok(vals
+            .iter()
+            .map(|value| value.into_string("", config))
+            .collect::<Vec<String>>()),
+        Value::List { vals, span } => vals
+            .iter()
+            .map(|value| extract_headers(value, config))
+            .next()
+            .ok_or_else(|| {
+                ShellError::SpannedLabeledError(
+                    "Found empty list".to_string(),
+                    "unable to extract headers".to_string(),
+                    *span,
+                )
+            })?,
+        _ => Err(ShellError::TypeMismatch(
+            "record".to_string(),
+            value.span()?,
+        )),
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(Headers {})
+    }
+}

crates/nu-command/src/filters/mod.rs

@@ -16,6 +16,7 @@ mod first;
 mod flatten;
 mod get;
 mod group_by;
+mod headers;
 mod keep;
 mod last;
 mod length;
@@ -62,6 +63,7 @@ pub use first::First;
 pub use flatten::Flatten;
 pub use get::Get;
 pub use group_by::GroupBy;
+pub use headers::Headers;
 pub use keep::*;
 pub use last::Last;
 pub use length::Length;