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}