From 666fbbb0d12af5d28e9e66891e9f7a505ff49441 Mon Sep 17 00:00:00 2001
From: Chris Gillespie <6572184+gillespiecd@users.noreply.github.com>
Date: Wed, 28 Oct 2020 20:14:08 -0700
Subject: [PATCH] Precision added to round cmd (#2710)

---
 crates/nu-cli/src/commands/math/round.rs | 88 +++++++++++++++---------
 crates/nu-cli/src/stream/input.rs        |  2 +-
 2 files changed, 58 insertions(+), 32 deletions(-)

diff --git a/crates/nu-cli/src/commands/math/round.rs b/crates/nu-cli/src/commands/math/round.rs
index f1a92dcd28..92ea30c783 100644
--- a/crates/nu-cli/src/commands/math/round.rs
+++ b/crates/nu-cli/src/commands/math/round.rs
@@ -1,11 +1,16 @@
-use crate::commands::math::utils::run_with_numerical_functions_on_stream;
 use crate::commands::WholeStreamCommand;
 use crate::prelude::*;
 use nu_errors::ShellError;
-use nu_protocol::{Signature, UntaggedValue, Value};
+use nu_protocol::{Primitive, Signature, SyntaxShape, UntaggedValue, Value};
+use nu_source::Tagged;
 
 pub struct SubCommand;
 
+#[derive(Deserialize)]
+struct Arguments {
+    precision: Option<Tagged<i64>>,
+}
+
 #[async_trait]
 impl WholeStreamCommand for SubCommand {
     fn name(&self) -> &str {
@@ -13,7 +18,12 @@ impl WholeStreamCommand for SubCommand {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build("math round")
+        Signature::build("math round").named(
+            "precision",
+            SyntaxShape::Number,
+            "digits of precision",
+            Some('p'),
+        )
     }
 
     fn usage(&self) -> &str {
@@ -25,44 +35,59 @@ impl WholeStreamCommand for SubCommand {
         args: CommandArgs,
         registry: &CommandRegistry,
     ) -> Result<OutputStream, ShellError> {
-        run_with_numerical_functions_on_stream(
-            RunnableContext {
-                input: args.input,
-                registry: registry.clone(),
-                shell_manager: args.shell_manager,
-                host: args.host,
-                ctrl_c: args.ctrl_c,
-                current_errors: args.current_errors,
-                name: args.call_info.name_tag,
-                raw_input: args.raw_input,
-            },
-            round_big_int,
-            round_big_decimal,
-            round_default,
-        )
-        .await
+        operate(args, registry).await
     }
 
     fn examples(&self) -> Vec<Example> {
-        vec![Example {
-            description: "Apply the round function to a list of numbers",
-            example: "echo [1.5 2.3 -3.1] | math round",
-            result: Some(vec![
-                UntaggedValue::int(2).into(),
-                UntaggedValue::int(2).into(),
-                UntaggedValue::int(-3).into(),
-            ]),
-        }]
+        vec![
+            Example {
+                description: "Apply the round function to a list of numbers",
+                example: "echo [1.5 2.3 -3.1] | math round",
+                result: Some(vec![
+                    UntaggedValue::int(2).into(),
+                    UntaggedValue::int(2).into(),
+                    UntaggedValue::int(-3).into(),
+                ]),
+            },
+            Example {
+                description: "Apply the round function with precision specified",
+                example: "echo [1.555 2.333 -3.111] | math round -p 2",
+                result: Some(vec![
+                    UntaggedValue::decimal_from_float(1.56, Span::default()).into(),
+                    UntaggedValue::decimal_from_float(2.33, Span::default()).into(),
+                    UntaggedValue::decimal_from_float(-3.11, Span::default()).into(),
+                ]),
+            },
+        ]
     }
 }
 
+async fn operate(
+    args: CommandArgs,
+    registry: &CommandRegistry,
+) -> Result<OutputStream, ShellError> {
+    let (Arguments { precision }, input) = args.process(&registry).await?;
+    let precision = precision.map(|p| p.item).unwrap_or(0);
+
+    let mapped = input.map(move |val| match val.value {
+        UntaggedValue::Primitive(Primitive::Int(val)) => round_big_int(val),
+        UntaggedValue::Primitive(Primitive::Decimal(val)) => round_big_decimal(val, precision),
+        other => round_default(other),
+    });
+    Ok(OutputStream::from_input(mapped))
+}
+
 fn round_big_int(val: BigInt) -> Value {
     UntaggedValue::int(val).into()
 }
 
-fn round_big_decimal(val: BigDecimal) -> Value {
-    let (rounded, _) = val.round(0).as_bigint_and_exponent();
-    UntaggedValue::int(rounded).into()
+fn round_big_decimal(val: BigDecimal, precision: i64) -> Value {
+    if precision > 0 {
+        UntaggedValue::decimal(val.round(precision)).into()
+    } else {
+        let (rounded, _) = val.round(precision).as_bigint_and_exponent();
+        UntaggedValue::int(rounded).into()
+    }
 }
 
 fn round_default(_: UntaggedValue) -> Value {
@@ -71,6 +96,7 @@ fn round_default(_: UntaggedValue) -> Value {
     ))
     .into()
 }
+
 #[cfg(test)]
 mod tests {
     use super::ShellError;
diff --git a/crates/nu-cli/src/stream/input.rs b/crates/nu-cli/src/stream/input.rs
index 280e0ff2d9..5137d633e3 100644
--- a/crates/nu-cli/src/stream/input.rs
+++ b/crates/nu-cli/src/stream/input.rs
@@ -7,7 +7,7 @@ use nu_source::{Tagged, TaggedItem};
 pub struct InputStream {
     values: BoxStream<'static, Value>,
 
-    // Whether or not an empty stream was explicitly requeted via InputStream::empty
+    // Whether or not an empty stream was explicitly requested via InputStream::empty
     empty: bool,
 }