Merge pull request #260 from luccasmmg/engine-q-math-2

Added math avg

crates/nu-command/src/default_context.rs

@@ -51,6 +51,7 @@ pub fn create_default_context() -> EngineState {
             Ls,
             Math,
             MathAbs,
+            MathAvg,
             Mkdir,
             Module,
             Mv,

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

@@ -0,0 +1,83 @@
+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 avg"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("math avg")
+    }
+
+    fn usage(&self) -> &str {
+        "Finds the average 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, average)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Get the average of a list of numbers",
+            example: "[-50 100.0 25] | math avg",
+            result: Some(Value::Float {
+                val: 25.0,
+                span: Span::unknown(),
+            }),
+        }]
+    }
+}
+
+pub fn average(values: &[Value], head: &Span) -> Result<Value, ShellError> {
+    let sum = reducer_for(Reduce::Summation);
+    let total = &sum(
+        Value::Int {
+            val: 0,
+            span: Span::unknown(),
+        },
+        values.to_vec(),
+    )?;
+    match total {
+        Value::Filesize { val, span } => Ok(Value::Filesize {
+            val: val / values.len() as i64,
+            span: *span,
+        }),
+        Value::Duration { val, span } => Ok(Value::Duration {
+            val: val / values.len() as i64,
+            span: *span,
+        }),
+        _ => total.div(
+            *head,
+            &Value::Int {
+                val: values.len() as i64,
+                span: *head,
+            },
+        ),
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(SubCommand {})
+    }
+}

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

@@ -1,5 +1,9 @@
 mod abs;
+mod avg;
 pub mod command;
+mod reducers;
+mod utils;
 
 pub use abs::SubCommand as MathAbs;
+pub use avg::SubCommand as MathAvg;
 pub use command::MathCommand as Math;

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

@@ -0,0 +1,60 @@
+use nu_protocol::{ShellError, Span, Value};
+
+#[allow(dead_code)]
+pub enum Reduce {
+    Summation,
+}
+
+pub fn reducer_for(
+    command: Reduce,
+) -> Box<dyn Fn(Value, Vec<Value>) -> Result<Value, ShellError> + Send + Sync + 'static> {
+    match command {
+        Reduce::Summation => Box::new(|_, values| sum(values)),
+    }
+}
+
+pub fn sum(data: Vec<Value>) -> Result<Value, ShellError> {
+    let initial_value = data.get(0);
+
+    let mut acc = match initial_value {
+        Some(Value::Filesize { span, .. }) => Ok(Value::Filesize {
+            val: 0,
+            span: *span,
+        }),
+        Some(Value::Duration { span, .. }) => Ok(Value::Duration {
+            val: 0,
+            span: *span,
+        }),
+        Some(Value::Int { span, .. }) | Some(Value::Float { span, .. }) => Ok(Value::Int {
+            val: 0,
+            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 { .. }
+            | Value::Filesize { .. }
+            | Value::Duration { .. } => {
+                let new_value = acc.add(acc.span().unwrap_or_else(|_| Span::unknown()), 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 sum of a value that cannot be summed".to_string(),
+                    other.span().unwrap_or_else(|_| Span::unknown()),
+                ));
+            }
+        }
+    }
+    Ok(acc)
+}

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

@@ -0,0 +1,87 @@
+use nu_protocol::ast::Call;
+use nu_protocol::{IntoPipelineData, PipelineData, ShellError, Span, Value};
+use std::collections::HashMap;
+
+pub type MathFunction = fn(values: &[Value], span: &Span) -> Result<Value, ShellError>;
+
+pub fn run_with_function(
+    call: &Call,
+    input: PipelineData,
+    mf: MathFunction,
+) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+    let name = call.head;
+    let res = calculate(input, name, mf);
+    match res {
+        Ok(v) => Ok(v.into_pipeline_data()),
+        Err(e) => Err(e),
+    }
+}
+
+fn helper_for_tables(
+    values: PipelineData,
+    name: Span,
+    mf: MathFunction,
+) -> Result<Value, ShellError> {
+    // If we are not dealing with Primitives, then perhaps we are dealing with a table
+    // Create a key for each column name
+    let mut column_values = HashMap::new();
+    for val in values {
+        if let Value::Record { cols, vals, .. } = val {
+            for (key, value) in cols.iter().zip(vals.iter()) {
+                column_values
+                    .entry(key.clone())
+                    .and_modify(|v: &mut Vec<Value>| v.push(value.clone()))
+                    .or_insert_with(|| vec![value.clone()]);
+            }
+        }
+    }
+    // The mathematical function operates over the columns of the table
+    let mut column_totals = HashMap::new();
+    for (col_name, col_vals) in column_values {
+        if let Ok(out) = mf(&col_vals, &name) {
+            column_totals.insert(col_name, out);
+        }
+    }
+    if column_totals.keys().len() == 0 {
+        return Err(ShellError::UnsupportedInput(
+            "Unable to give a result with this input".to_string(),
+            name,
+        ));
+    }
+    let (cols, vals) = column_totals
+        .into_iter()
+        .fold((vec![], vec![]), |mut acc, (k, v)| {
+            acc.0.push(k);
+            acc.1.push(v);
+            acc
+        });
+
+    Ok(Value::Record {
+        cols,
+        vals,
+        span: name,
+    })
+}
+
+pub fn calculate(values: PipelineData, name: Span, mf: MathFunction) -> Result<Value, ShellError> {
+    match values {
+        PipelineData::Stream(_) => helper_for_tables(values, name, mf),
+        PipelineData::Value(Value::List { ref vals, .. }) => match &vals[..] {
+            [Value::Record { .. }, _end @ ..] => helper_for_tables(values, name, mf),
+            _ => mf(vals, &name),
+        },
+        PipelineData::Value(Value::Record { vals, cols, span }) => {
+            let new_vals: Result<Vec<Value>, ShellError> =
+                vals.into_iter().map(|val| mf(&[val], &name)).collect();
+            match new_vals {
+                Ok(vec) => Ok(Value::Record {
+                    cols,
+                    vals: vec,
+                    span,
+                }),
+                Err(err) => Err(err),
+            }
+        }
+        PipelineData::Value(val) => mf(&[val], &name),
+    }
+}