From 411435d68f9318218f3ffbade1ba90619db2bfcc Mon Sep 17 00:00:00 2001
From: Fernando Herrera <fernando.j.herrera@gmail.com>
Date: Tue, 20 Jul 2021 13:07:42 +0100
Subject: [PATCH] Dataframe Shape command (#3805)

* size command to get dataframe info

* change command name to shape

* apply lint to file
---
 .../nu-command/src/commands/dataframe/mod.rs  |  2 +
 .../src/commands/dataframe/shape.rs           | 47 +++++++++++++++++++
 crates/nu-command/src/commands/mod.rs         |  7 +--
 crates/nu-command/src/default_context.rs      |  1 +
 4 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 crates/nu-command/src/commands/dataframe/shape.rs

diff --git a/crates/nu-command/src/commands/dataframe/mod.rs b/crates/nu-command/src/commands/dataframe/mod.rs
index 48f33da6a7..85100460e6 100644
--- a/crates/nu-command/src/commands/dataframe/mod.rs
+++ b/crates/nu-command/src/commands/dataframe/mod.rs
@@ -18,6 +18,7 @@ pub mod open;
 pub mod pivot;
 pub mod sample;
 pub mod select;
+pub mod shape;
 pub mod show;
 pub mod slice;
 pub mod sort;
@@ -50,6 +51,7 @@ pub use open::DataFrame as DataFrameOpen;
 pub use pivot::DataFrame as DataFramePivot;
 pub use sample::DataFrame as DataFrameSample;
 pub use select::DataFrame as DataFrameSelect;
+pub use shape::DataFrame as DataFrameShape;
 pub use show::DataFrame as DataFrameShow;
 pub use slice::DataFrame as DataFrameSlice;
 pub use sort::DataFrame as DataFrameSort;
diff --git a/crates/nu-command/src/commands/dataframe/shape.rs b/crates/nu-command/src/commands/dataframe/shape.rs
new file mode 100644
index 0000000000..111c6f56bf
--- /dev/null
+++ b/crates/nu-command/src/commands/dataframe/shape.rs
@@ -0,0 +1,47 @@
+use crate::prelude::*;
+use nu_engine::WholeStreamCommand;
+use nu_errors::ShellError;
+use nu_protocol::{dataframe::NuDataFrame, Signature, TaggedDictBuilder};
+
+pub struct DataFrame;
+
+impl WholeStreamCommand for DataFrame {
+    fn name(&self) -> &str {
+        "dataframe shape"
+    }
+
+    fn usage(&self) -> &str {
+        "[DataFrame] Shows column and row size for a dataframe"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("dataframe shape")
+    }
+
+    fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
+        command(args)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Shows row and column shape",
+            example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe shape",
+            result: None,
+        }]
+    }
+}
+
+fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
+    let tag = args.call_info.name_tag.clone();
+
+    let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
+
+    let rows = df.as_ref().height();
+    let cols = df.as_ref().width();
+
+    let mut data = TaggedDictBuilder::new(&tag);
+    data.insert_value("rows", format!("{}", rows));
+    data.insert_value("columns", format!("{}", cols));
+
+    Ok(OutputStream::one(data.into_value()))
+}
diff --git a/crates/nu-command/src/commands/mod.rs b/crates/nu-command/src/commands/mod.rs
index 46acd65d38..f738d5f11e 100644
--- a/crates/nu-command/src/commands/mod.rs
+++ b/crates/nu-command/src/commands/mod.rs
@@ -32,9 +32,10 @@ pub use dataframe::{
     DataFrameIsIn, DataFrameIsNotNull, DataFrameIsNull, DataFrameIsUnique, DataFrameJoin,
     DataFrameLast, DataFrameList, DataFrameMelt, DataFrameNNull, DataFrameNUnique, DataFrameNot,
     DataFrameOpen, DataFramePivot, DataFrameSample, DataFrameSelect, DataFrameSeriesRename,
-    DataFrameSet, DataFrameSetWithIdx, DataFrameShift, DataFrameShow, DataFrameSlice,
-    DataFrameSort, DataFrameTake, DataFrameToCsv, DataFrameToDF, DataFrameToParquet,
-    DataFrameToSeries, DataFrameUnique, DataFrameValueCounts, DataFrameWhere, DataFrameWithColumn,
+    DataFrameSet, DataFrameSetWithIdx, DataFrameShape, DataFrameShift, DataFrameShow,
+    DataFrameSlice, DataFrameSort, DataFrameTake, DataFrameToCsv, DataFrameToDF,
+    DataFrameToParquet, DataFrameToSeries, DataFrameUnique, DataFrameValueCounts, DataFrameWhere,
+    DataFrameWithColumn,
 };
 pub use env::*;
 pub use filesystem::*;
diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs
index 8a591d449d..2f82656d6e 100644
--- a/crates/nu-command/src/default_context.rs
+++ b/crates/nu-command/src/default_context.rs
@@ -317,6 +317,7 @@ pub fn create_default_context(interactive: bool) -> Result<EvaluationContext, Bo
             whole_stream_command(DataFrameNot),
             whole_stream_command(DataFrameTake),
             whole_stream_command(DataFrameSetWithIdx),
+            whole_stream_command(DataFrameShape),
         ]);
 
         #[cfg(feature = "clipboard-cli")]