1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

use super::variance::compute_variance as variance;
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{Dictionary, Primitive, Signature, UntaggedValue, Value};
use std::str::FromStr;

pub struct SubCommand;

impl WholeStreamCommand for SubCommand {
    fn name(&self) -> &str {
        "math stddev"
    }

    fn signature(&self) -> Signature {
        Signature::build("math stddev").switch(
            "sample",
            "calculate sample standard deviation",
            Some('s'),
        )
    }

    fn usage(&self) -> &str {
        "Finds the stddev of a list of numbers or tables"
    }

    fn run(&self, mut args: CommandArgs) -> Result<OutputStream, ShellError> {
        let sample: bool = args.has_flag("sample");
        let values: Vec<Value> = args.input.drain_vec();
        let name = args.call_info.name_tag;

        let n = if sample {
            values.len() - 1
        } else {
            values.len()
        };

        let res = if values.iter().all(|v| v.is_primitive()) {
            compute_stddev(&values, n, &name)
        } else {
            // 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 = IndexMap::new();
            for value in values {
                if let UntaggedValue::Row(row_dict) = &value.value {
                    for (key, value) in &row_dict.entries {
                        column_values
                            .entry(key.clone())
                            .and_modify(|v: &mut Vec<Value>| v.push(value.clone()))
                            .or_insert(vec![value.clone()]);
                    }
                }
            }
            // The mathematical function operates over the columns of the table
            let mut column_totals = IndexMap::new();
            for (col_name, col_vals) in column_values {
                if let Ok(out) = compute_stddev(&col_vals, n, &name) {
                    column_totals.insert(col_name, out);
                }
            }

            if column_totals.keys().len() == 0 {
                return Err(ShellError::labeled_error(
                    "Attempted to compute values that can't be operated on",
                    "value appears here",
                    name.span,
                ));
            }

            Ok(UntaggedValue::Row(Dictionary {
                entries: column_totals,
            })
            .into_untagged_value())
        }?;

        if res.value.is_table() {
            Ok(OutputStream::from(
                res.table_entries().cloned().collect::<Vec<_>>(),
            ))
        } else {
            Ok(OutputStream::one(res))
        }
    }

    fn examples(&self) -> Vec<Example> {
        vec![
            Example {
                description: "Get the stddev of a list of numbers",
                example: "echo [1 2 3 4 5] | math stddev",
                result: Some(vec![UntaggedValue::decimal(BigDecimal::from_str("1.414213562373095048801688724209698078569671875376948073176679737990732478462107038850387534327641573").expect("Could not convert to decimal from string")).into()]),
            },
            Example {
                description: "Get the sample stddev of a list of numbers",
                example: "echo [1 2 3 4 5] | math stddev -s",
                result: Some(vec![UntaggedValue::decimal(BigDecimal::from_str("1.581138830084189665999446772216359266859777569662608413428752426396297219319619110672124054189650148").expect("Could not convert to decimal from string")).into()]),
            },
        ]
    }
}

#[cfg(test)]
pub fn stddev(values: &[Value], name: &Tag) -> Result<Value, ShellError> {
    compute_stddev(values, values.len(), name)
}

pub fn compute_stddev(values: &[Value], n: usize, name: &Tag) -> Result<Value, ShellError> {
    let variance = variance(values, n, name)?.as_primitive()?;
    let sqrt_var = match variance {
        Primitive::Decimal(var) => var.sqrt(),
        _ => {
            return Err(ShellError::labeled_error(
                "Could not take square root of variance",
                "error occurred here",
                name.span,
            ))
        }
    };
    match sqrt_var {
        Some(stddev) => Ok(UntaggedValue::from(Primitive::Decimal(stddev)).into_value(name)),
        None => Err(ShellError::labeled_error(
            "Could not calculate stddev",
            "error occurred here",
            name.span,
        )),
    }
}

#[cfg(test)]
mod tests {
    use super::ShellError;
    use super::SubCommand;

    #[test]
    fn examples_work_as_expected() -> Result<(), ShellError> {
        use crate::examples::test as test_examples;

        test_examples(SubCommand {})
    }
}