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 tag = args.call_info.name_tag.clone();
-        let rows: Vec<Value> = args.input.collect().await;
+        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 {
+        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

docs/commands/count.md

@@ -1,6 +1,10 @@
 # count
 
-This command counts the number of rows in a table.
+Obtain the row or column count of a table.
+
+## Flags
+
+* `-c`, `--column`: Calculate number of columns in table
 
 ## Examples
 
@@ -32,14 +36,18 @@ This command counts the number of rows in a table.
 ────┴────────────────────┴──────┴──────────┴──────────────
 ```
 
+By default, `count` will return the number of rows in a table
+
 ```shell
 > ls | count
 20
 ```
 
+The `-c` flag will produce a count of the columns in the table
+
 ```shell
-> ls | get name | count
+> ls | count -c
-20
+4
 ```
 
 ```shell