From 5d58f68c599021c3401f628f511b6f0cc2ded38b Mon Sep 17 00:00:00 2001
From: Michael Angerman <stormasm@gmail.com>
Date: Fri, 31 Dec 2021 20:27:20 -0800
Subject: [PATCH] port over from nushell the column flag for the length command
 (#617)

* port over from nushell the column flag for the length command

* fix clippy error

* refactor with the get_columns now centrally located
---
 crates/nu-command/src/filters/length.rs | 94 +++++++++++++++++++++----
 src/tests/test_table_operations.rs      | 13 ++++
 2 files changed, 92 insertions(+), 15 deletions(-)

diff --git a/crates/nu-command/src/filters/length.rs b/crates/nu-command/src/filters/length.rs
index bddf977b9..d28637e9e 100644
--- a/crates/nu-command/src/filters/length.rs
+++ b/crates/nu-command/src/filters/length.rs
@@ -1,6 +1,10 @@
+use nu_engine::column::get_columns;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
-use nu_protocol::{Category, IntoPipelineData, PipelineData, Signature, Value};
+use nu_protocol::{
+    Category, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, ShellError, Signature,
+    Span, Value,
+};
 
 #[derive(Clone)]
 pub struct Length;
@@ -15,27 +19,87 @@ impl Command for Length {
     }
 
     fn signature(&self) -> nu_protocol::Signature {
-        Signature::build("length").category(Category::Filters)
+        Signature::build("length")
+            .switch("column", "Show the number of columns in a table", Some('c'))
+            .category(Category::Filters)
     }
 
     fn run(
         &self,
-        _engine_state: &EngineState,
+        engine_state: &EngineState,
         _stack: &mut Stack,
         call: &Call,
         input: PipelineData,
-    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
-        match input {
-            PipelineData::Value(Value::Nothing { .. }, ..) => Ok(Value::Int {
-                val: 0,
-                span: call.head,
-            }
-            .into_pipeline_data()),
-            _ => Ok(Value::Int {
-                val: input.into_iter().count() as i64,
-                span: call.head,
-            }
-            .into_pipeline_data()),
+    ) -> Result<nu_protocol::PipelineData, ShellError> {
+        let col = call.has_flag("column");
+        if col {
+            length_col(engine_state, call, input)
+        } else {
+            length_row(call, input)
+        }
+    }
+}
+
+// this simulates calling input | columns | length
+fn length_col(
+    engine_state: &EngineState,
+    call: &Call,
+    input: PipelineData,
+) -> Result<PipelineData, ShellError> {
+    length_row(
+        call,
+        getcol(engine_state, call.head, input)
+            .expect("getcol() should not fail used in column command"),
+    )
+}
+
+fn length_row(call: &Call, input: PipelineData) -> Result<PipelineData, ShellError> {
+    match input {
+        PipelineData::Value(Value::Nothing { .. }, ..) => Ok(Value::Int {
+            val: 0,
+            span: call.head,
+        }
+        .into_pipeline_data()),
+        _ => Ok(Value::Int {
+            val: input.into_iter().count() as i64,
+            span: call.head,
+        }
+        .into_pipeline_data()),
+    }
+}
+
+fn getcol(
+    engine_state: &EngineState,
+    span: Span,
+    input: PipelineData,
+) -> Result<PipelineData, ShellError> {
+    match input {
+        PipelineData::Value(
+            Value::List {
+                vals: input_vals,
+                span,
+            },
+            ..,
+        ) => {
+            let input_cols = get_columns(&input_vals);
+            Ok(input_cols
+                .into_iter()
+                .map(move |x| Value::String { val: x, span })
+                .into_pipeline_data(engine_state.ctrlc.clone()))
+        }
+        PipelineData::ListStream(stream, ..) => {
+            let v: Vec<_> = stream.into_iter().collect();
+            let input_cols = get_columns(&v);
+
+            Ok(input_cols
+                .into_iter()
+                .map(move |x| Value::String { val: x, span })
+                .into_pipeline_data(engine_state.ctrlc.clone()))
+        }
+        PipelineData::Value(..) | PipelineData::StringStream(..) | PipelineData::ByteStream(..) => {
+            let cols = vec![];
+            let vals = vec![];
+            Ok(Value::Record { cols, vals, span }.into_pipeline_data())
         }
     }
 }
diff --git a/src/tests/test_table_operations.rs b/src/tests/test_table_operations.rs
index 8ce6f7355..6ccb2528d 100644
--- a/src/tests/test_table_operations.rs
+++ b/src/tests/test_table_operations.rs
@@ -197,3 +197,16 @@ fn select() -> TestResult {
 fn update_will_insert() -> TestResult {
     run_test(r#"{} | update a b | get a"#, "b")
 }
+
+#[test]
+fn length_for_columns() -> TestResult {
+    run_test(
+        r#"[[name,age,grade]; [bill,20,a] [a b c]] | length -c"#,
+        "3",
+    )
+}
+
+#[test]
+fn length_for_rows() -> TestResult {
+    run_test(r#"[[name,age,grade]; [bill,20,a] [a b c]] | length"#, "2")
+}