nu_command/math/
avg.rs

1use crate::math::{
2    reducers::{Reduce, reducer_for},
3    utils::run_with_function,
4};
5use nu_engine::command_prelude::*;
6
7const NS_PER_SEC: i64 = 1_000_000_000;
8#[derive(Clone)]
9pub struct MathAvg;
10
11impl Command for MathAvg {
12    fn name(&self) -> &str {
13        "math avg"
14    }
15
16    fn signature(&self) -> Signature {
17        Signature::build("math avg")
18            .input_output_types(vec![
19                (Type::List(Box::new(Type::Duration)), Type::Duration),
20                (Type::Duration, Type::Duration),
21                (Type::List(Box::new(Type::Filesize)), Type::Filesize),
22                (Type::Filesize, Type::Filesize),
23                (Type::List(Box::new(Type::Number)), Type::Number),
24                (Type::Number, Type::Number),
25                (Type::Range, Type::Number),
26                (Type::table(), Type::record()),
27                (Type::record(), Type::record()),
28            ])
29            .allow_variants_without_examples(true)
30            .category(Category::Math)
31    }
32
33    fn description(&self) -> &str {
34        "Returns the average of a list of numbers."
35    }
36
37    fn search_terms(&self) -> Vec<&str> {
38        vec!["average", "mean", "statistics"]
39    }
40
41    fn is_const(&self) -> bool {
42        true
43    }
44
45    fn run(
46        &self,
47        _engine_state: &EngineState,
48        _stack: &mut Stack,
49        call: &Call,
50        input: PipelineData,
51    ) -> Result<PipelineData, ShellError> {
52        run_with_function(call, input, average)
53    }
54
55    fn run_const(
56        &self,
57        _working_set: &StateWorkingSet,
58        call: &Call,
59        input: PipelineData,
60    ) -> Result<PipelineData, ShellError> {
61        run_with_function(call, input, average)
62    }
63
64    fn examples(&self) -> Vec<Example> {
65        vec![
66            Example {
67                description: "Compute the average of a list of numbers",
68                example: "[-50 100.0 25] | math avg",
69                result: Some(Value::test_float(25.0)),
70            },
71            Example {
72                description: "Compute the average of a list of durations",
73                example: "[2sec 1min] | math avg",
74                result: Some(Value::test_duration(31 * NS_PER_SEC)),
75            },
76            Example {
77                description: "Compute the average of each column in a table",
78                example: "[[a b]; [1 2] [3 4]] | math avg",
79                result: Some(Value::test_record(record! {
80                    "a" => Value::test_int(2),
81                    "b" => Value::test_int(3),
82                })),
83            },
84        ]
85    }
86}
87
88pub fn average(values: &[Value], span: Span, head: Span) -> Result<Value, ShellError> {
89    let sum = reducer_for(Reduce::Summation);
90    let total = &sum(Value::int(0, head), values.to_vec(), span, head)?;
91    let span = total.span();
92    match total {
93        Value::Filesize { val, .. } => Ok(Value::filesize(val.get() / values.len() as i64, span)),
94        Value::Duration { val, .. } => Ok(Value::duration(val / values.len() as i64, span)),
95        _ => total.div(head, &Value::int(values.len() as i64, head), head),
96    }
97}
98
99#[cfg(test)]
100mod test {
101    use super::*;
102
103    #[test]
104    fn test_examples() {
105        use crate::test_examples;
106
107        test_examples(MathAvg {})
108    }
109}