Add mean() and median()

expr-lang · Aug 30, 2023 · 3b2df5a · 3b2df5a
1 parent ed3d1f3
commit 3b2df5a
Show file tree

Hide file tree

Showing 2 changed files with 143 additions and 0 deletions.
diff --git a/builtin/builtin.go b/builtin/builtin.go
@@ -367,6 +367,136 @@ var Builtins = []*ast.Function{
 			return args[0], nil
 		},
 	},
+	{
+		Name: "sum",
+		Func: func(args ...any) (any, error) {
+			if len(args) != 1 {
+				return nil, fmt.Errorf("invalid number of arguments (expected 1, got %d)", len(args))
+			}
+			v := reflect.ValueOf(args[0])
+			if v.Kind() != reflect.Slice && v.Kind() != reflect.Array {
+				return nil, fmt.Errorf("cannot sum %s", v.Kind())
+			}
+			sum := int64(0)
+			i := 0
+			for ; i < v.Len(); i++ {
+				it := deref(v.Index(i))
+				if it.CanInt() {
+					sum += it.Int()
+				} else if it.CanFloat() {
+					goto float
+				} else {
+					return nil, fmt.Errorf("cannot sum %s", it.Kind())
+				}
+			}
+			return int(sum), nil
+		float:
+			fSum := float64(sum)
+			for ; i < v.Len(); i++ {
+				it := deref(v.Index(i))
+				if it.CanInt() {
+					fSum += float64(it.Int())
+				} else if it.CanFloat() {
+					fSum += it.Float()
+				} else {
+					return nil, fmt.Errorf("cannot sum %s", it.Kind())
+				}
+			}
+			return fSum, nil
+		},
+		Validate: func(args []reflect.Type) (reflect.Type, error) {
+			if len(args) != 1 {
+				return anyType, fmt.Errorf("invalid number of arguments (expected 1, got %d)", len(args))
+			}
+			switch kind(args[0]) {
+			case reflect.Interface, reflect.Slice, reflect.Array:
+			default:
+				return anyType, fmt.Errorf("cannot sum %s", args[0])
+			}
+			return anyType, nil
+		},
+	},
+	{
+		Name: "mean",
+		Func: func(args ...any) (any, error) {
+			if len(args) != 1 {
+				return nil, fmt.Errorf("invalid number of arguments (expected 1, got %d)", len(args))
+			}
+			v := reflect.ValueOf(args[0])
+			if v.Kind() != reflect.Slice && v.Kind() != reflect.Array {
+				return nil, fmt.Errorf("cannot mean %s", v.Kind())
+			}
+			if v.Len() == 0 {
+				return 0.0, nil
+			}
+			sum := float64(0)
+			i := 0
+			for ; i < v.Len(); i++ {
+				it := deref(v.Index(i))
+				if it.CanInt() {
+					sum += float64(it.Int())
+				} else if it.CanFloat() {
+					sum += it.Float()
+				} else {
+					return nil, fmt.Errorf("cannot mean %s", it.Kind())
+				}
+			}
+			return sum / float64(i), nil
+		},
+		Validate: func(args []reflect.Type) (reflect.Type, error) {
+			if len(args) != 1 {
+				return anyType, fmt.Errorf("invalid number of arguments (expected 1, got %d)", len(args))
+			}
+			switch kind(args[0]) {
+			case reflect.Interface, reflect.Slice, reflect.Array:
+			default:
+				return anyType, fmt.Errorf("cannot avg %s", args[0])
+			}
+			return floatType, nil
+		},
+	},
+	{
+		Name: "median",
+		Func: func(args ...any) (any, error) {
+			if len(args) != 1 {
+				return nil, fmt.Errorf("invalid number of arguments (expected 1, got %d)", len(args))
+			}
+			v := reflect.ValueOf(args[0])
+			if v.Kind() != reflect.Slice && v.Kind() != reflect.Array {
+				return nil, fmt.Errorf("cannot median %s", v.Kind())
+			}
+			if v.Len() == 0 {
+				return 0.0, nil
+			}
+			s := make([]float64, v.Len())
+			for i := 0; i < v.Len(); i++ {
+				it := deref(v.Index(i))
+				if it.CanInt() {
+					s[i] = float64(it.Int())
+				} else if it.CanFloat() {
+					s[i] = it.Float()
+				} else {
+					return nil, fmt.Errorf("cannot median %s", it.Kind())
+				}
+			}
+			sort.Float64s(s)
+			if len(s)%2 == 0 {
+				return (s[len(s)/2-1] + s[len(s)/2]) / 2, nil
+			}
+			return s[len(s)/2], nil
+		},
+		Validate: func(args []reflect.Type) (reflect.Type, error) {
+			if len(args) != 1 {
+				return anyType, fmt.Errorf("invalid number of arguments (expected 1, got %d)", len(args))
+			}
+			switch kind(args[0]) {
+			case reflect.Interface, reflect.Slice, reflect.Array:
+			default:
+				return anyType, fmt.Errorf("cannot median %s", args[0])
+			}
+			return floatType, nil
+		},
+	},
 	{
 		Name: "toJSON",
 		Func: func(args ...any) (any, error) {

diff --git a/builtin/builtin_test.go b/builtin/builtin_test.go
@@ -69,6 +69,19 @@ func TestBuiltin(t *testing.T) {
 		{`max(1.5, 2.5, 3.5)`, 3.5},
 		{`min(1, 2, 3)`, 1},
 		{`min(1.5, 2.5, 3.5)`, 1.5},
+		{`sum(1..9)`, 45},
+		{`sum([.5, 1.5, 2.5])`, 4.5},
+		{`sum([])`, 0},
+		{`sum([1, 2, 3.0, 4])`, 10.0},
+		{`mean(1..9)`, 5.0},
+		{`mean([.5, 1.5, 2.5])`, 1.5},
+		{`mean([])`, 0.0},
+		{`mean([1, 2, 3.0, 4])`, 2.5},
+		{`median(1..9)`, 5.0},
+		{`median([.5, 1.5, 2.5])`, 1.5},
+		{`median([])`, 0.0},
+		{`median([1, 2, 3])`, 2.0},
+		{`median([1, 2, 3, 4])`, 2.5},
 		{`toJSON({foo: 1, bar: 2})`, "{\n  \"bar\": 2,\n  \"foo\": 1\n}"},
 		{`fromJSON("[1, 2, 3]")`, []any{1.0, 2.0, 3.0}},
 		{`toBase64("hello")`, "aGVsbG8="},