time: per-thread time.now() function

lib/time/time.go

@@ -71,8 +71,25 @@ var Module = &starlarkstruct.Module{
 // NowFunc is a function that generates 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: use SetNow() to set per-thread clock.
 var NowFunc = time.Now
 
+const contextKey = "time.Now"
+
+// SetNow associates with the specified Starlark thread a function that generates the current time.
+// Can be used for example by applications that require their
+// Starlark scripts to be fully deterministic.
+func SetNow(thread *starlark.Thread, nowFunc func() time.Time) {
+	thread.SetLocal(contextKey, nowFunc)
+}
+
+// Now returns the function to generate the current time previously associated with this thread.
+func Now(thread *starlark.Thread) func() time.Time {
+	nowFunc, _ := thread.Local(contextKey).(func() time.Time)
+	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 +146,11 @@ 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
+	nowFunc := Now(thread)
+	if nowFunc == nil {
+		nowFunc = NowFunc
+	}
+	return Time(nowFunc()), nil
 }
 
 // Duration is a Starlark representation of a duration.