diff --git a/crates/nu-cli/src/cli.rs b/crates/nu-cli/src/cli.rs index 308ad6499..7b633e0f7 100644 --- a/crates/nu-cli/src/cli.rs +++ b/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 diff --git a/crates/nu-cli/src/commands.rs b/crates/nu-cli/src/commands.rs index 8345bf227..942df8773 100644 --- a/crates/nu-cli/src/commands.rs +++ b/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; diff --git a/crates/nu-cli/src/commands/benchmark.rs b/crates/nu-cli/src/commands/benchmark.rs new file mode 100644 index 000000000..84bd3ab4a --- /dev/null +++ b/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 { + benchmark(args, registry).await + } +} + +async fn benchmark( + raw_args: CommandArgs, + registry: &CommandRegistry, +) -> Result { + let registry = registry.clone(); + + let mut context = Context::from_raw(&raw_args, ®istry); + let scope = raw_args.call_info.scope.clone(); + let (BenchmarkArgs { block }, input) = raw_args.process(®istry).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])) +}