time: per-thread time.now() function

lib/time/time.go

@@ -18,36 +18,36 @@ import (
 // Module time is a Starlark module of time-related functions and constants.
 // The module defines the following functions:
 //
-//     from_timestamp(sec, nsec) - Converts the given Unix time corresponding to the number of seconds
-//                                 and (optionally) nanoseconds since January 1, 1970 UTC into an object
-//                                 of type Time. For more details, refer to https://pkg.go.dev/time#Unix.
+//	    from_timestamp(sec, nsec) - Converts the given Unix time corresponding to the number of seconds
+//	                                and (optionally) nanoseconds since January 1, 1970 UTC into an object
+//	                                of type Time. For more details, refer to https://pkg.go.dev/time#Unix.
 //
-//     is_valid_timezone(loc) - Reports whether loc is a valid time zone name.
+//	    is_valid_timezone(loc) - Reports whether loc is a valid time zone name.
 //
-//     now() - Returns the current local time. Applications may replace this function by a deterministic one.
+//	    now() - Returns the current local time. Applications may replace this function by a deterministic one.
 //
-//     parse_duration(d) - Parses the given duration string. For more details, refer to
-//                         https://pkg.go.dev/time#ParseDuration.
+//	    parse_duration(d) - Parses the given duration string. For more details, refer to
+//	                        https://pkg.go.dev/time#ParseDuration.
 //
-//     parse_time(x, format, location) - Parses the given time string using a specific time format and location.
-//                                      The expected arguments are a time string (mandatory), a time format
-//                                      (optional, set to RFC3339 by default, e.g. "2021-03-22T23:20:50.52Z")
-//                                      and a name of location (optional, set to UTC by default). For more details,
-//                                      refer to https://pkg.go.dev/time#Parse and https://pkg.go.dev/time#ParseInLocation.
+//	    parse_time(x, format, location) - Parses the given time string using a specific time format and location.
+//	                                     The expected arguments are a time string (mandatory), a time format
+//	                                     (optional, set to RFC3339 by default, e.g. "2021-03-22T23:20:50.52Z")
+//	                                     and a name of location (optional, set to UTC by default). For more details,
+//	                                     refer to https://pkg.go.dev/time#Parse and https://pkg.go.dev/time#ParseInLocation.
 //
-//     time(year, month, day, hour, minute, second, nanosecond, location) - Returns the Time corresponding to
-//	                                                                        yyyy-mm-dd hh:mm:ss + nsec nanoseconds
-//                                                                          in the appropriate zone for that time
-//                                                                          in the given location. All the parameters
-//                                                                          are optional.
+//	    time(year, month, day, hour, minute, second, nanosecond, location) - Returns the Time corresponding to
+//		                                                                        yyyy-mm-dd hh:mm:ss + nsec nanoseconds
+//	                                                                         in the appropriate zone for that time
+//	                                                                         in the given location. All the parameters
+//	                                                                         are optional.
 // The module also defines the following constants:
 //
-//     nanosecond - A duration representing one nanosecond.
-//     microsecond - A duration representing one microsecond.
-//     millisecond - A duration representing one millisecond.
-//     second - A duration representing one second.
-//     minute - A duration representing one minute.
-//     hour - A duration representing one hour.
+//	nanosecond - A duration representing one nanosecond.
+//	microsecond - A duration representing one microsecond.
+//	millisecond - A duration representing one millisecond.
+//	second - A duration representing one second.
+//	minute - A duration representing one minute.
+//	hour - A duration representing one hour.
 //
 var Module = &starlarkstruct.Module{
 	Name: "time",
@@ -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.
@@ -222,14 +243,14 @@ func (d Duration) Cmp(v starlark.Value, depth int) (int, error) {
 
 // Binary implements binary operators, which satisfies the starlark.HasBinary
 // interface. operators:
-//    duration + duration = duration
-//    duration + time = time
-//    duration - duration = duration
-//    duration / duration = float
-//    duration / int = duration
-//    duration / float = duration
-//    duration // duration = int
-//    duration * int = duration
+//	duration + duration = duration
+//	duration + time = time
+//	duration - duration = duration
+//	duration / duration = float
+//	duration / int = duration
+//	duration / float = duration
+//	duration // duration = int
+//	duration * int = duration
 func (d Duration) Binary(op syntax.Token, y starlark.Value, side starlark.Side) (starlark.Value, error) {
 	x := time.Duration(d)
 
@@ -406,9 +427,9 @@ func (t Time) Cmp(yV starlark.Value, depth int) (int, error) {
 
 // Binary implements binary operators, which satisfies the starlark.HasBinary
 // interface
-//    time + duration = time
-//    time - duration = time
-//    time - time = duration
+//	time + duration = time
+//	time - duration = time
+//	time - time = duration
 func (t Time) Binary(op syntax.Token, y starlark.Value, side starlark.Side) (starlark.Value, error) {
 	x := time.Time(t)