Skip to main content

nu_command/math/
avg.rs

1use crate::math::{
2    reducers::{Reduce, reducer_for},
3    utils::{run_with_function_with_cell_paths, run_with_function_with_cell_paths_const},
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            .rest(
31                "columns",
32                SyntaxShape::CellPath,
33                "The cell-paths/columns to operate on.",
34            )
35            .category(Category::Math)
36    }
37
38    fn description(&self) -> &str {
39        "Returns the average of a list of numbers."
40    }
41
42    fn search_terms(&self) -> Vec<&str> {
43        vec!["average", "mean", "statistics"]
44    }
45
46    fn is_const(&self) -> bool {
47        true
48    }
49
50    fn run(
51        &self,
52        engine_state: &EngineState,
53        stack: &mut Stack,
54        call: &Call,
55        input: PipelineData,
56    ) -> Result<PipelineData, ShellError> {
57        run_with_function_with_cell_paths(engine_state, stack, call, input, average)
58    }
59
60    fn run_const(
61        &self,
62        working_set: &StateWorkingSet,
63        call: &Call,
64        input: PipelineData,
65    ) -> Result<PipelineData, ShellError> {
66        run_with_function_with_cell_paths_const(working_set, call, input, average)
67    }
68
69    fn examples(&self) -> Vec<Example<'_>> {
70        vec![
71            Example {
72                description: "Compute the average of a list of numbers.",
73                example: "[-50 100.0 25] | math avg",
74                result: Some(Value::test_float(25.0)),
75            },
76            Example {
77                description: "Compute the average of a list of durations.",
78                example: "[2sec 1min] | math avg",
79                result: Some(Value::test_duration(31 * NS_PER_SEC)),
80            },
81            Example {
82                description: "Compute the average of each column in a table.",
83                example: "[[a b]; [1 2] [3 4]] | math avg",
84                result: Some(Value::test_record(record! {
85                    "a" => Value::test_int(2),
86                    "b" => Value::test_int(3),
87                })),
88            },
89            Example {
90                description: "Compute the average of list-valued columns in a record.",
91                example: "{alice: [1 2 3], bob: [4 5 6]} | math avg",
92                result: Some(Value::test_record(record! {
93                    "alice" => Value::test_int(2),
94                    "bob" => Value::test_int(5),
95                })),
96            },
97            Example {
98                description: "Compute the average of a single column using a cell path.",
99                example: "{alice: [1 2 3], bob: [4 5 6]} | math avg alice",
100                result: Some(Value::test_record(record! {
101                    "alice" => Value::test_int(2),
102                    "bob" => Value::list(
103                        vec![Value::test_int(4), Value::test_int(5), Value::test_int(6)],
104                        Span::test_data(),
105                    ),
106                })),
107            },
108        ]
109    }
110}
111
112pub fn average(values: &[Value], span: Span, head: Span) -> Result<Value, ShellError> {
113    let sum = reducer_for(Reduce::Summation);
114    let total = &sum(Value::int(0, head), values.to_vec(), span, head)?;
115    let span = total.span();
116    match total {
117        Value::Filesize { val, .. } => Ok(Value::filesize(val.get() / values.len() as i64, span)),
118        Value::Duration { val, .. } => Ok(Value::duration(val / values.len() as i64, span)),
119        _ => total.div(head, &Value::int(values.len() as i64, head), head),
120    }
121}
122
123#[cfg(test)]
124mod test {
125    use super::*;
126
127    #[test]
128    fn test_examples() -> nu_test_support::Result {
129        nu_test_support::test().examples(MathAvg)
130    }
131}