From 40ad9acbc36ce5077632462c109ebac1afad846d Mon Sep 17 00:00:00 2001 From: Luccas Mateus de Medeiros Gomes Date: Tue, 26 Oct 2021 22:00:50 -0300 Subject: [PATCH] Added math avg Linting Fix clippy warning Fix list of records --- crates/nu-command/src/default_context.rs | 1 + crates/nu-command/src/math/avg.rs | 83 ++++++++++++++++++++++ crates/nu-command/src/math/mod.rs | 4 ++ crates/nu-command/src/math/reducers.rs | 60 ++++++++++++++++ crates/nu-command/src/math/utils.rs | 87 ++++++++++++++++++++++++ 5 files changed, 235 insertions(+) create mode 100644 crates/nu-command/src/math/avg.rs create mode 100644 crates/nu-command/src/math/reducers.rs create mode 100644 crates/nu-command/src/math/utils.rs diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 64b68f19dd..b1f8fd929c 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -50,6 +50,7 @@ pub fn create_default_context() -> EngineState { Ls, Math, MathAbs, + MathAvg, Mkdir, Module, Mv, diff --git a/crates/nu-command/src/math/avg.rs b/crates/nu-command/src/math/avg.rs new file mode 100644 index 0000000000..384d3493a5 --- /dev/null +++ b/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 { + run_with_function(call, input, average) + } + + fn examples(&self) -> Vec { + 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 { + 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 {}) + } +} diff --git a/crates/nu-command/src/math/mod.rs b/crates/nu-command/src/math/mod.rs index 5abb051b23..c3f4a42934 100644 --- a/crates/nu-command/src/math/mod.rs +++ b/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; diff --git a/crates/nu-command/src/math/reducers.rs b/crates/nu-command/src/math/reducers.rs new file mode 100644 index 0000000000..2339fd97f8 --- /dev/null +++ b/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) -> Result + Send + Sync + 'static> { + match command { + Reduce::Summation => Box::new(|_, values| sum(values)), + } +} + +pub fn sum(data: Vec) -> Result { + 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) +} diff --git a/crates/nu-command/src/math/utils.rs b/crates/nu-command/src/math/utils.rs new file mode 100644 index 0000000000..770940c4c0 --- /dev/null +++ b/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; + +pub fn run_with_function( + call: &Call, + input: PipelineData, + mf: MathFunction, +) -> Result { + 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 { + // 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| 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 { + 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, 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), + } +}