Add column flag to count command

crates/nu-cli/src/commands/count.rs

@@ -7,6 +7,11 @@ use nu_protocol::{Signature, UntaggedValue, Value};
 
 pub struct Count;
 
+#[derive(Deserialize)]
+pub struct CountArgs {
+    column: bool,
+}
+
 #[async_trait]
 impl WholeStreamCommand for Count {
     fn name(&self) -> &str {
@@ -14,7 +19,11 @@ impl WholeStreamCommand for Count {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build("count")
+        Signature::build("count").switch(
+            "column",
+            "Calculate number of columns in table",
+            Some('c'),
+        )
     }
 
     fn usage(&self) -> &str {
@@ -24,22 +33,44 @@ impl WholeStreamCommand for Count {
     async fn run(
         &self,
         args: CommandArgs,
-        _registry: &CommandRegistry,
+        registry: &CommandRegistry,
     ) -> Result<OutputStream, ShellError> {
-        let name = args.call_info.name_tag.clone();
-        let rows: Vec<Value> = args.input.collect().await;
+        let tag = args.call_info.name_tag.clone();
+        let (CountArgs { column }, input) = args.process(&registry).await?;
+        let rows: Vec<Value> = input.collect().await;
+
+        if column {
+            if let UntaggedValue::Row(dict) = &rows[0].value {
+                return Ok(OutputStream::one(
+                    UntaggedValue::int(dict.length()).into_value(tag),
+                ));
+            } else {
+                return Err(ShellError::labeled_error(
+                    "Cannot obtain column count",
+                    "cannot obtain column count",
+                    tag,
+                ));
+            }
+        }
 
         Ok(OutputStream::one(
-            UntaggedValue::int(rows.len()).into_value(name),
+            UntaggedValue::int(rows.len()).into_value(tag),
         ))
     }
 
     fn examples(&self) -> Vec<Example> {
-        vec![Example {
-            description: "Count the number of entries in a list",
-            example: "echo [1 2 3 4 5] | count",
-            result: Some(vec![UntaggedValue::int(5).into()]),
-        }]
+        vec![
+            Example {
+                description: "Count the number of entries in a list",
+                example: "echo [1 2 3 4 5] | count",
+                result: Some(vec![UntaggedValue::int(5).into()]),
+            },
+            Example {
+                description: "Count the number of columns in the calendar table",
+                example: "cal | count -c",
+                result: None,
+            },
+        ]
     }
 }
 

crates/nu-protocol/src/value/dict.rs

@@ -188,6 +188,11 @@ impl Dictionary {
     pub fn insert_data_at_key(&mut self, name: &str, value: Value) {
         self.entries.insert(name.to_string(), value);
     }
+
+    /// Return size of dictionary
+    pub fn length(&self) -> usize {
+        self.entries.len()
+    }
 }
 
 /// A helper to help create dictionaries for you. It has the ability to insert values into the dictionary while maintaining the tags that need to be applied to the individual members