add math module

cmd/starlark/starlark.go

@@ -16,6 +16,7 @@ import (
 	"strings"
 
 	"go.starlark.net/internal/compile"
+	"go.starlark.net/lib/math"
 	"go.starlark.net/lib/time"
 	"go.starlark.net/repl"
 	"go.starlark.net/resolve"
@@ -93,6 +94,7 @@ func doMain() int {
 	// TODO(adonovan): plumb predeclared env through to the REPL.
 	starlark.Universe["json"] = starlarkjson.Module
 	starlark.Universe["time"] = time.Module
+	starlark.Universe["math"] = math.Module
 
 	switch {
 	case flag.NArg() == 1 || *execprog != "":

lib/math/math.go

@@ -0,0 +1,150 @@
+// Copyright 2021 The Bazel Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package math provides basic constants and mathematical functions.
+package math // import "go.starlark.net/lib/math"
+
+import (
+	"fmt"
+	"math"
+
+	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
+)
+
+// Module math is a Starlark module of math-related functions and constants.
+// The module defines the following functions:
+//
+//     abs(x) - Returns the absolute value of x.
+//     ceil(x) - Returns the ceiling of x, the smallest integer greater than or equal to x.
+//     floor(x) - Returns the floor of x, the largest integer less than or equal to x.
+//     round(x) - Returns the nearest integer, rounding half away from zero.
+//
+//     exp(x) - Returns e raised to the power x, where e = 2.718281… is the base of natural logarithms.
+//     sqrt(x) - Returns the square root of x.
+//
+//     acos(x) - Returns the arc cosine of x, in radians.
+//     asin(x) - Returns the arc sine of x, in radians.
+//     atan(x) - Returns the arc tangent of x, in radians.
+//     atan2(y, x) - Returns atan(y / x), in radians.
+//                   The result is between -pi and pi.
+//                   The vector in the plane from the origin to point (x, y) makes this angle with the positive X axis.
+//                   The point of atan2() is that the signs of both inputs are known to it, so it can compute the correct
+//                   quadrant for the angle.
+//                   For example, atan(1) and atan2(1, 1) are both pi/4, but atan2(-1, -1) is -3*pi/4.
+//     cos(x) - Returns the cosine of x, in radians.
+//     hypot(x, y) - Returns the Euclidean norm, sqrt(x*x + y*y). This is the length of the vector from the origin to point (x, y).
+//     sin(x) - Returns the sine of x, in radians.
+//     tan(x) - Returns the tangent of x, in radians.
+//
+//     degrees(x) - Converts angle x from radians to degrees.
+//     radians(x) - Converts angle x from degrees to radians.
+//
+//     acosh(x) - Returns the inverse hyperbolic cosine of x.
+//     asinh(x) - Returns the inverse hyperbolic sine of x.
+//     atanh(x) - Returns the inverse hyperbolic tangent of x.
+//     cosh(x) - Returns the hyperbolic cosine of x.
+//     sinh(x) - Returns the hyperbolic sine of x.
+//     tanh(x) - Returns the hyperbolic tangent of x.
+//
+//     log(x) - Returns the natural logarithm of x.
+//     log2(x) - Returns the decimal logarithm of x. The special cases are the same as for Log.
+//     log10(x) - Returns the natural logarithm of 1 plus its argument x. It is more accurate than Log(1 + x) when x is near.
+//
+// All function accept both int and float values as arguments.
+//
+// The module also defines approximations of the following constants:
+//
+//     e - The base of natural logarithms, approximately 2.71828.
+//     phi - The golden ratio, (1 + sqrt(5)) / 2, approximately 1.61803.
+//     pi - The ratio of a circle's circumference to its diameter, approximately 3.14159.
+//
+var Module = &starlarkstruct.Module{
+	Name: "math",
+	Members: starlark.StringDict{
+		"abs":   newUnaryBuiltin("abs", math.Abs),
+		"ceil":  newUnaryBuiltin("ceil", math.Ceil),
+		"floor": newUnaryBuiltin("floor", math.Floor),
+		"round": newUnaryBuiltin("round", math.Round),
+
+		"exp":  newUnaryBuiltin("exp", math.Exp),
+		"sqrt": newUnaryBuiltin("sqrt", math.Sqrt),
+
+		"acos":  newUnaryBuiltin("acos", math.Acos),
+		"asin":  newUnaryBuiltin("asin", math.Asin),
+		"atan":  newUnaryBuiltin("atan", math.Atan),
+		"atan2": newBinaryFunction("atan2", math.Atan2),
+		"cos":   newUnaryBuiltin("cos", math.Cos),
+		"hypot": newBinaryFunction("hypot", math.Hypot),
+		"sin":   newUnaryBuiltin("sin", math.Sin),
+		"tan":   newUnaryBuiltin("tan", math.Tan),
+
+		"degrees": newUnaryBuiltin("degrees", degrees),
+		"radians": newUnaryBuiltin("radians", radians),
+
+		"acosh": newUnaryBuiltin("acosh", math.Acosh),
+		"asinh": newUnaryBuiltin("asinh", math.Asinh),
+		"atanh": newUnaryBuiltin("atanh", math.Atanh),
+		"cosh":  newUnaryBuiltin("cosh", math.Cosh),
+		"sinh":  newUnaryBuiltin("sinh", math.Sinh),
+		"tanh":  newUnaryBuiltin("tanh", math.Tanh),
+
+		"log":   newUnaryBuiltin("log", math.Log),
+		"log2":  newUnaryBuiltin("log2", math.Log2),
+		"log10": newUnaryBuiltin("log10", math.Log10),
+
+		"e":   starlark.Float(math.E),
+		"phi": starlark.Float(math.Phi),
+		"pi":  starlark.Float(math.Pi),
+	},
+}
+
+// floatOrInt is an Unpacker that converts a Starlark int or float to Go's float64.
+type floatOrInt float64
+
+// Unpack is a custom argument unpacker
+func (p *floatOrInt) Unpack(v starlark.Value) error {
+	switch v := v.(type) {
+	case starlark.Int:
+		*p = floatOrInt(v.Float())
+		return nil
+	case starlark.Float:
+		*p = floatOrInt(v)
+		return nil
+	}
+
+	return fmt.Errorf("got %s, want float or int", v.Type())
+}
+
+// newUnaryBuiltin wraps a unary floating-point Go function
+// as a Starlark built-in that accepts int or float arguments.
+func newUnaryBuiltin(name string, fn func(float64) float64) *starlark.Builtin {
+	return starlark.NewBuiltin(name, func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+		var x floatOrInt
+		if err := starlark.UnpackPositionalArgs(name, args, kwargs, 1, &x); err != nil {
+			return nil, err
+		}
+		return starlark.Float(fn(float64(x))), nil
+	})
+}
+
+// newBinaryFunction wraps a binary floating-point Go function
+// as a Starlark built-in that accepts int or float arguments.
+func newBinaryFunction(name string, fn func(float64, float64) float64) *starlark.Builtin {
+	return starlark.NewBuiltin(name, func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+		var x, y floatOrInt
+		if err := starlark.UnpackPositionalArgs(name, args, kwargs, 2, &x, &y); err != nil {
+			return nil, err
+		}
+		return starlark.Float(fn(float64(x), float64(y))), nil
+	})
+}
+
+func degrees(x float64) float64 {
+	return 360 * x / (2 * math.Pi)
+}
+
+func radians(x float64) float64 {
+	return 2 * math.Pi * x / 360
+}

starlark/eval_test.go

@@ -16,6 +16,7 @@ import (
 	"testing"
 
 	"go.starlark.net/internal/chunkedfile"
+	starlarkmath "go.starlark.net/lib/math"
 	"go.starlark.net/lib/time"
 	"go.starlark.net/resolve"
 	"go.starlark.net/starlark"
@@ -124,6 +125,7 @@ func TestExecFile(t *testing.T) {
 		"testdata/int.star",
 		"testdata/json.star",
 		"testdata/list.star",
+		"testdata/math.star",
 		"testdata/misc.star",
 		"testdata/set.star",
 		"testdata/string.star",
@@ -198,6 +200,9 @@ func load(thread *starlark.Thread, module string) (starlark.StringDict, error) {
 	if module == "time.star" {
 		return starlark.StringDict{"time": time.Module}, nil
 	}
+	if module == "math.star" {
+		return starlark.StringDict{"math": starlarkmath.Module}, nil
+	}
 
 	// TODO(adonovan): test load() using this execution path.
 	filename := filepath.Join(filepath.Dir(thread.CallFrame(0).Pos.Filename()), module)