Add benchmark command (#2223)

crates/nu-cli/src/cli.rs

@@ -278,6 +278,7 @@ pub fn create_default_context(
             // Statistics
             whole_stream_command(Size),
             whole_stream_command(Count),
+            whole_stream_command(Benchmark),
             // Metadata
             whole_stream_command(Tags),
             // Shells

crates/nu-cli/src/commands.rs

@@ -12,6 +12,7 @@ pub(crate) mod autoenv;
 pub(crate) mod autoenv_trust;
 pub(crate) mod autoenv_untrust;
 pub(crate) mod autoview;
+pub(crate) mod benchmark;
 pub(crate) mod build_string;
 pub(crate) mod cal;
 pub(crate) mod cd;
@@ -131,6 +132,7 @@ pub(crate) use append::Append;
 pub(crate) use autoenv::Autoenv;
 pub(crate) use autoenv_trust::AutoenvTrust;
 pub(crate) use autoenv_untrust::AutoenvUnTrust;
+pub(crate) use benchmark::Benchmark;
 pub(crate) use build_string::BuildString;
 pub(crate) use cal::Cal;
 pub(crate) use char_::Char;

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

@@ -0,0 +1,80 @@
+use crate::commands::classified::block::run_block;
+use crate::commands::WholeStreamCommand;
+use crate::prelude::*;
+use nu_errors::ShellError;
+use nu_protocol::{
+    hir::Block, Primitive, ReturnSuccess, Signature, SyntaxShape, UntaggedValue, Value,
+};
+
+use chrono::prelude::*;
+
+pub struct Benchmark;
+
+#[derive(Deserialize, Debug)]
+struct BenchmarkArgs {
+    block: Block,
+}
+
+#[async_trait]
+impl WholeStreamCommand for Benchmark {
+    fn name(&self) -> &str {
+        "benchmark"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("benchmark").required(
+            "block",
+            SyntaxShape::Block,
+            "the block to run and benchmark",
+        )
+    }
+
+    fn usage(&self) -> &str {
+        "Runs a block and return the time it took to do execute it. Eg) benchmark { echo $nu.env.NAME }"
+    }
+
+    async fn run(
+        &self,
+        args: CommandArgs,
+        registry: &CommandRegistry,
+    ) -> Result<OutputStream, ShellError> {
+        benchmark(args, registry).await
+    }
+}
+
+async fn benchmark(
+    raw_args: CommandArgs,
+    registry: &CommandRegistry,
+) -> Result<OutputStream, ShellError> {
+    let registry = registry.clone();
+
+    let mut context = Context::from_raw(&raw_args, &registry);
+    let scope = raw_args.call_info.scope.clone();
+    let (BenchmarkArgs { block }, input) = raw_args.process(&registry).await?;
+
+    let start_time: chrono::DateTime<_> = Utc::now();
+
+    let result = run_block(
+        &block,
+        &mut context,
+        input,
+        &scope.it,
+        &scope.vars,
+        &scope.env,
+    )
+    .await;
+
+    let output = match result {
+        Ok(mut stream) => {
+            let _ = stream.drain_vec().await;
+            let run_duration: chrono::Duration = Utc::now().signed_duration_since(start_time);
+            context.clear_errors();
+            Ok(ReturnSuccess::Value(Value {
+                value: UntaggedValue::Primitive(Primitive::from(run_duration)),
+                tag: Tag::from(block.span),
+            }))
+        }
+        Err(e) => Err(e),
+    };
+    Ok(OutputStream::from(vec![output]))
+}