From 45f32c9541615bd44d35b633ddadaf62a6b044f2 Mon Sep 17 00:00:00 2001
From: gorogoroumaru <zokutyou2@gmail.com>
Date: Sun, 13 Sep 2020 08:56:05 +0900
Subject: [PATCH] Allow `math avg` to work on durations (#2529)

* Allow `math avg` to work on durations

* formatting

* fix linting issue and implemented `math sum` for duration

* fix linting issue

* applied requested changes

* applied requested change for avg.rs

* formatting
---
 crates/nu-cli/src/commands/math/avg.rs | 22 ++++++++++++++++++++++
 crates/nu-cli/src/commands/math/sum.rs |  1 +
 crates/nu-data/src/value.rs            | 23 +++++++++++++++++++++++
 crates/nu-protocol/src/value.rs        |  4 ++++
 4 files changed, 50 insertions(+)

diff --git a/crates/nu-cli/src/commands/math/avg.rs b/crates/nu-cli/src/commands/math/avg.rs
index 5f89620913..c24841fb9a 100644
--- a/crates/nu-cli/src/commands/math/avg.rs
+++ b/crates/nu-cli/src/commands/math/avg.rs
@@ -137,6 +137,28 @@ pub fn average(values: &[Value], name: &Tag) -> Result<Value, ShellError> {
                 )),
             }
         }
+        Value {
+            value: UntaggedValue::Primitive(Primitive::Duration(duration)),
+            ..
+        } => {
+            let left = UntaggedValue::from(Primitive::Duration(duration));
+            let result = nu_data::value::compute_values(Operator::Divide, &left, &total_rows);
+
+            match result {
+                Ok(UntaggedValue::Primitive(Primitive::Duration(result))) => {
+                    Ok(UntaggedValue::duration(result).into_value(name))
+                }
+                Ok(_) => Err(ShellError::labeled_error(
+                    "could not calculate average of non-integer or unrelated types",
+                    "source",
+                    name,
+                )),
+                Err((left_type, right_type)) => Err(ShellError::coerce_error(
+                    left_type.spanned(name.span),
+                    right_type.spanned(name.span),
+                )),
+            }
+        }
         Value {
             value: UntaggedValue::Primitive(other),
             ..
diff --git a/crates/nu-cli/src/commands/math/sum.rs b/crates/nu-cli/src/commands/math/sum.rs
index a5d4aecd43..6005be35f7 100644
--- a/crates/nu-cli/src/commands/math/sum.rs
+++ b/crates/nu-cli/src/commands/math/sum.rs
@@ -103,6 +103,7 @@ pub fn summation(values: &[Value], name: &Tag) -> Result<Value, ShellError> {
                 &name.span,
             )
         }),
+        v if v.is_duration() => sum(UntaggedValue::int(0).into_untagged_value(), values.to_vec()),
         // v is nothing primitive
         v if v.is_none() => sum(
             UntaggedValue::int(0).into_untagged_value(),
diff --git a/crates/nu-data/src/value.rs b/crates/nu-data/src/value.rs
index eecc2e9fa1..40ac374e2c 100644
--- a/crates/nu-data/src/value.rs
+++ b/crates/nu-data/src/value.rs
@@ -212,6 +212,29 @@ pub fn compute_values(
 
                 Ok(UntaggedValue::Primitive(Primitive::Duration(result)))
             }
+            (Primitive::Int(x), Primitive::Duration(y)) => {
+                let result = match operator {
+                    Operator::Plus => Ok(x + y),
+                    Operator::Minus => Ok(x - y),
+                    _ => Err((left.type_name(), right.type_name())),
+                }?;
+
+                Ok(UntaggedValue::Primitive(Primitive::Duration(result)))
+            }
+            (Primitive::Duration(x), Primitive::Decimal(y)) => {
+                let result = match operator {
+                    Operator::Divide => {
+                        if y.is_zero() {
+                            return Ok(zero_division_error());
+                        }
+                        let y = y.as_bigint_and_exponent();
+                        Ok(x / y.0)
+                    }
+                    _ => Err((left.type_name(), right.type_name())),
+                }?;
+
+                Ok(UntaggedValue::Primitive(Primitive::Duration(result)))
+            }
             _ => Err((left.type_name(), right.type_name())),
         },
         _ => Err((left.type_name(), right.type_name())),
diff --git a/crates/nu-protocol/src/value.rs b/crates/nu-protocol/src/value.rs
index 2daa937c29..eeb22ac774 100644
--- a/crates/nu-protocol/src/value.rs
+++ b/crates/nu-protocol/src/value.rs
@@ -83,6 +83,10 @@ impl UntaggedValue {
         matches!(self, UntaggedValue::Primitive(Primitive::Filesize(_)))
     }
 
+    pub fn is_duration(&self) -> bool {
+        matches!(self, UntaggedValue::Primitive(Primitive::Duration(_)))
+    }
+
     /// Returns true if this value represents a table
     pub fn is_table(&self) -> bool {
         matches!(self, UntaggedValue::Table(_))