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}