Merge pull request #273 from luccasmmg/engine-q-math

New math commands(product, round, sqrt and sum)

crates/nu-command/src/default_context.rs

@@ -63,6 +63,10 @@ pub fn create_default_context() -> EngineState {
             MathAvg,
             MathMax,
             MathMin,
+            MathProduct,
+            MathRound,
+            MathSqrt,
+            MathSum,
             Mkdir,
             Module,
             Mv,

crates/nu-command/src/math/mod.rs

@@ -3,7 +3,11 @@ mod avg;
 pub mod command;
 mod max;
 mod min;
+mod product;
 mod reducers;
+mod round;
+mod sqrt;
+mod sum;
 mod utils;
 
 pub use abs::SubCommand as MathAbs;
@@ -11,3 +15,7 @@ pub use avg::SubCommand as MathAvg;
 pub use command::MathCommand as Math;
 pub use max::SubCommand as MathMax;
 pub use min::SubCommand as MathMin;
+pub use product::SubCommand as MathProduct;
+pub use round::SubCommand as MathRound;
+pub use sqrt::SubCommand as MathSqrt;
+pub use sum::SubCommand as MathSum;

crates/nu-command/src/math/product.rs

@@ -0,0 +1,58 @@
+use crate::math::reducers::{reducer_for, Reduce};
+use crate::math::utils::run_with_function;
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{Example, PipelineData, ShellError, Signature, Span, Value};
+
+#[derive(Clone)]
+pub struct SubCommand;
+
+impl Command for SubCommand {
+    fn name(&self) -> &str {
+        "math product"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("math product")
+    }
+
+    fn usage(&self) -> &str {
+        "Finds the product of a list of numbers or tables"
+    }
+
+    fn run(
+        &self,
+        _engine_state: &EngineState,
+        _stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        run_with_function(call, input, product)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Get the product of a list of numbers",
+            example: "[2 3 3 4] | math product",
+            result: Some(Value::test_int(72)),
+        }]
+    }
+}
+
+/// Calculate product of given values
+pub fn product(values: &[Value], head: &Span) -> Result<Value, ShellError> {
+    let product_func = reducer_for(Reduce::Product);
+    product_func(Value::nothing(), values.to_vec(), *head)
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(SubCommand {})
+    }
+}

crates/nu-command/src/math/reducers.rs

@@ -4,6 +4,7 @@ use std::cmp::Ordering;
 #[allow(dead_code)]
 pub enum Reduce {
     Summation,
+    Product,
     Minimum,
     Maximum,
 }
@@ -14,6 +15,7 @@ pub type ReducerFunction =
 pub fn reducer_for(command: Reduce) -> ReducerFunction {
     match command {
         Reduce::Summation => Box::new(|_, values, head| sum(values, head)),
+        Reduce::Product => Box::new(|_, values, head| product(values, head)),
         Reduce::Minimum => Box::new(|_, values, head| min(values, head)),
         Reduce::Maximum => Box::new(|_, values, head| max(values, head)),
     }
@@ -112,3 +114,39 @@ pub fn sum(data: Vec<Value>, head: Span) -> Result<Value, ShellError> {
     }
     Ok(acc)
 }
+
+pub fn product(data: Vec<Value>, head: Span) -> Result<Value, ShellError> {
+    let initial_value = data.get(0);
+
+    let mut acc = match initial_value {
+        Some(Value::Int { span, .. }) | Some(Value::Float { span, .. }) => Ok(Value::Int {
+            val: 1,
+            span: *span,
+        }),
+        None => Err(ShellError::UnsupportedInput(
+            "Empty input".to_string(),
+            Span::unknown(),
+        )),
+        _ => Ok(Value::nothing()),
+    }?;
+
+    for value in &data {
+        match value {
+            Value::Int { .. } | Value::Float { .. } => {
+                let new_value = acc.mul(head, value);
+                if new_value.is_err() {
+                    return new_value;
+                }
+                acc = new_value.expect("This should never trigger")
+            }
+            other => {
+                return Err(ShellError::UnsupportedInput(
+                    "Attempted to compute the product of a value that cannot be multiplied"
+                        .to_string(),
+                    other.span().unwrap_or_else(|_| Span::unknown()),
+                ));
+            }
+        }
+    }
+    Ok(acc)
+}

crates/nu-command/src/math/round.rs

@@ -0,0 +1,110 @@
+use nu_engine::CallExt;
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value};
+
+#[derive(Clone)]
+pub struct SubCommand;
+
+impl Command for SubCommand {
+    fn name(&self) -> &str {
+        "math round"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("math round").named(
+            "precision",
+            SyntaxShape::Number,
+            "digits of precision",
+            Some('p'),
+        )
+    }
+
+    fn usage(&self) -> &str {
+        "Applies the round function to a list of numbers"
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        let precision_param: Option<i64> = call.get_flag(engine_state, stack, "precision")?;
+        let head = call.head;
+        input.map(
+            move |value| operate(value, head, precision_param),
+            engine_state.ctrlc.clone(),
+        )
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Apply the round function to a list of numbers",
+                example: "[1.5 2.3 -3.1] | math round",
+                result: Some(Value::List {
+                    vals: vec![Value::test_int(2), Value::test_int(2), Value::test_int(-3)],
+                    span: Span::unknown(),
+                }),
+            },
+            Example {
+                description: "Apply the round function with precision specified",
+                example: "[1.555 2.333 -3.111] | math round -p 2",
+                result: Some(Value::List {
+                    vals: vec![
+                        Value::Float {
+                            val: 1.56,
+                            span: Span::unknown(),
+                        },
+                        Value::Float {
+                            val: 2.33,
+                            span: Span::unknown(),
+                        },
+                        Value::Float {
+                            val: -3.11,
+                            span: Span::unknown(),
+                        },
+                    ],
+                    span: Span::unknown(),
+                }),
+            },
+        ]
+    }
+}
+
+fn operate(value: Value, head: Span, precision: Option<i64>) -> Value {
+    match value {
+        Value::Float { val, span } => match precision {
+            Some(precision_number) => Value::Float {
+                val: ((val * ((10_f64).powf(precision_number as f64))).round()
+                    / (10_f64).powf(precision_number as f64)),
+                span,
+            },
+            None => Value::Int {
+                val: val.round() as i64,
+                span,
+            },
+        },
+        Value::Int { .. } => value,
+        other => Value::Error {
+            error: ShellError::UnsupportedInput(
+                String::from("Only numerical values are supported"),
+                other.span().unwrap_or(head),
+            ),
+        },
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(SubCommand {})
+    }
+}

crates/nu-command/src/math/sqrt.rs

@@ -0,0 +1,91 @@
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{Example, PipelineData, ShellError, Signature, Span, Value};
+
+#[derive(Clone)]
+pub struct SubCommand;
+
+impl Command for SubCommand {
+    fn name(&self) -> &str {
+        "math sqrt"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("math sqrt")
+    }
+
+    fn usage(&self) -> &str {
+        "Applies the square root function to a list of numbers"
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        _stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        let head = call.head;
+        input.map(
+            move |value| operate(value, head),
+            engine_state.ctrlc.clone(),
+        )
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Apply the square root function to a list of numbers",
+            example: "[9 16] | math sqrt",
+            result: Some(Value::List {
+                vals: vec![Value::test_int(3), Value::test_int(4)],
+                span: Span::unknown(),
+            }),
+        }]
+    }
+}
+
+fn operate(value: Value, head: Span) -> Value {
+    match value {
+        Value::Int { val, span } => {
+            let squared = (val as f64).sqrt();
+            if squared.is_nan() {
+                return error_negative_sqrt(span);
+            }
+            Value::Float { val: squared, span }
+        }
+        Value::Float { val, span } => {
+            let squared = val.sqrt();
+            if squared.is_nan() {
+                return error_negative_sqrt(span);
+            }
+            Value::Float { val: squared, span }
+        }
+        other => Value::Error {
+            error: ShellError::UnsupportedInput(
+                String::from("Only numerical values are supported"),
+                other.span().unwrap_or(head),
+            ),
+        },
+    }
+}
+
+fn error_negative_sqrt(span: Span) -> Value {
+    Value::Error {
+        error: ShellError::UnsupportedInput(
+            String::from("Can't square root a negative number"),
+            span,
+        ),
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(SubCommand {})
+    }
+}

crates/nu-command/src/math/sum.rs

@@ -0,0 +1,64 @@
+use crate::math::reducers::{reducer_for, Reduce};
+use crate::math::utils::run_with_function;
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{Example, PipelineData, ShellError, Signature, Span, Value};
+
+#[derive(Clone)]
+pub struct SubCommand;
+
+impl Command for SubCommand {
+    fn name(&self) -> &str {
+        "math sum"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("math sum")
+    }
+
+    fn usage(&self) -> &str {
+        "Finds the sum of a list of numbers or tables"
+    }
+
+    fn run(
+        &self,
+        _engine_state: &EngineState,
+        _stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        run_with_function(call, input, summation)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Sum a list of numbers",
+                example: "[1 2 3] | math sum",
+                result: Some(Value::test_int(6)),
+            },
+            Example {
+                description: "Get the disk usage for the current directory",
+                example: "ls | get size | math sum",
+                result: None,
+            },
+        ]
+    }
+}
+
+pub fn summation(values: &[Value], head: &Span) -> Result<Value, ShellError> {
+    let sum_func = reducer_for(Reduce::Summation);
+    sum_func(Value::nothing(), values.to_vec(), *head)
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(SubCommand {})
+    }
+}