time: per-thread time.now() function

lib/time/time.go

@@ -6,6 +6,7 @@
 package time // import "go.starlark.net/lib/time"
 
 import (
+	"errors"
 	"fmt"
 	"sort"
 	"time"
@@ -68,11 +69,29 @@ var Module = &starlarkstruct.Module{
 	},
 }
 
-// NowFunc is a function that generates the current time. Intentionally exported
+// NowFunc is a function that reports the current time. Intentionally exported
 // so that it can be overridden, for example by applications that require their
 // Starlark scripts to be fully deterministic.
+//
+// Deprecated: avoid updating this global variable
+// and instead use SetNow on each thread to set its clock function.
 var NowFunc = time.Now
 
+const contextKey = "time.now"
+
+// SetNow sets the thread's optional clock function.
+// If non-nil, it will be used in preference to NowFunc when the
+// thread requests the current time by executing a call to time.now.
+func SetNow(thread *starlark.Thread, nowFunc func() (time.Time, error)) {
+	thread.SetLocal(contextKey, nowFunc)
+}
+
+// Now returns the clock function previously associated with this thread.
+func Now(thread *starlark.Thread) func() (time.Time,error) {
+	nowFunc, _ := thread.Local(contextKey).(func() (time.Time,error))
+	return nowFunc
+}
+
 func parseDuration(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var d Duration
 	err := starlark.UnpackPositionalArgs("parse_duration", args, kwargs, 1, &d)
@@ -129,7 +148,19 @@ func fromTimestamp(thread *starlark.Thread, _ *starlark.Builtin, args starlark.T
 }
 
 func now(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
-	return Time(NowFunc()), nil
+	nowErrFunc := Now(thread)
+	if nowErrFunc != nil {
+		t, err := nowErrFunc()
+		if err != nil {
+			return nil, err
+		}
+		return Time(t), nil
+	}
+	nowFunc := NowFunc
+	if nowFunc == nil {
+		return nil, errors.New("time.now() is not available")
+	}
+	return Time(nowFunc()), nil
 }
 
 // Duration is a Starlark representation of a duration.

lib/time/time_test.go

@@ -0,0 +1,82 @@
+package time
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"go.starlark.net/starlark"
+)
+
+func TestPerThreadNowReturnsCorrectTime(t *testing.T) {
+	th := &starlark.Thread{}
+	date := time.Date(1, 2, 3, 4, 5, 6, 7, time.UTC)
+	SetNow(th, func() (time.Time, error) {
+		return date, nil
+	})
+
+	res, err := starlark.Call(th, Module.Members["now"], nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	retTime := time.Time(res.(Time))
+
+	if !retTime.Equal(date) {
+		t.Fatal("Expected time to be equal", retTime, date)
+	}
+}
+
+func TestPerThreadNowReturnsError(t *testing.T) {
+	th := &starlark.Thread{}
+	e := errors.New("no time")
+	SetNow(th, func() (time.Time, error) {
+		return time.Time{}, e
+	})
+
+	_, err := starlark.Call(th, Module.Members["now"], nil, nil)
+	if !errors.Is(err, e) {
+		t.Fatal("Expected equal error", e, err)
+	}
+}
+
+func TestGlobalNowReturnsCorrectTime(t *testing.T) {
+	th := &starlark.Thread{}
+
+	oldNow := NowFunc
+	defer func() {
+		NowFunc = oldNow
+	}()
+
+	date := time.Date(1, 2, 3, 4, 5, 6, 7, time.UTC)
+	NowFunc = func() time.Time {
+		return date
+	}
+
+	res, err := starlark.Call(th, Module.Members["now"], nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	retTime := time.Time(res.(Time))
+
+	if !retTime.Equal(date) {
+		t.Fatal("Expected time to be equal", retTime, date)
+	}
+}
+
+func TestGlobalNowReturnsErrorWhenNil(t *testing.T) {
+	th := &starlark.Thread{}
+
+	oldNow := NowFunc
+	defer func() {
+		NowFunc = oldNow
+	}()
+
+	NowFunc = nil
+
+	_, err := starlark.Call(th, Module.Members["now"], nil, nil)
+	if err == nil {
+		t.Fatal("Expected to get an error")
+	}
+}