diff --git a/crates/nu-cli/src/commands/count.rs b/crates/nu-cli/src/commands/count.rs index 7d0e846e70..969b79f31e 100644 --- a/crates/nu-cli/src/commands/count.rs +++ b/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 { - let name = args.call_info.name_tag.clone(); - let rows: Vec = args.input.collect().await; + let tag = args.call_info.name_tag.clone(); + let (CountArgs { column }, input) = args.process(®istry).await?; + let rows: Vec = 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 { - 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, + }, + ] } } diff --git a/crates/nu-protocol/src/value/dict.rs b/crates/nu-protocol/src/value/dict.rs index 9a82d6cc7c..cd2ddf0878 100644 --- a/crates/nu-protocol/src/value/dict.rs +++ b/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 diff --git a/docs/commands/count.md b/docs/commands/count.md index 028b819b55..8f7bd5188f 100644 --- a/docs/commands/count.md +++ b/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 -20 +> ls | count -c +4 ``` ```shell