From 35ad986f940d2c4980ce963cf6f0ffb47203a5bf Mon Sep 17 00:00:00 2001 From: garvinp-stripe <97996462+garvinp-stripe@users.noreply.github.com> Date: Fri, 3 Feb 2023 11:49:48 -0800 Subject: [PATCH] Add flatten on eval (#104) * add flatten * update protobuf for zlib + comment * fix test * add flatten as an option * remove conditional --- WORKSPACE | 7 +++---- skycfg.go | 33 ++++++++++++++++++++++++------ skycfg_test.go | 55 +++++++++++++++++++++++++++++++++++++++++++------- 3 files changed, 78 insertions(+), 17 deletions(-) diff --git a/WORKSPACE b/WORKSPACE index b9f7efb..01519f9 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -34,11 +34,10 @@ http_archive( http_archive( name = "rules_proto", - sha256 = "602e7161d9195e50246177e7c55b2f39950a9cf7366f74ed5f22fd45750cd208", - strip_prefix = "rules_proto-97d8af4dc474595af3900dd85cb3a29ad28cc313", + sha256 = "9850fcf6ad40fa348e6f13b2cfef4bb4639762f804794f2bf61d988f4ba0dae9", + strip_prefix = "rules_proto-4.0.0-3.19.2-2", urls = [ - "https://mirror.bazel.build/github.com/bazelbuild/rules_proto/archive/97d8af4dc474595af3900dd85cb3a29ad28cc313.tar.gz", - "https://github.com/bazelbuild/rules_proto/archive/97d8af4dc474595af3900dd85cb3a29ad28cc313.tar.gz", + "https://github.com/bazelbuild/rules_proto/archive/refs/tags/4.0.0-3.19.2-2.tar.gz", ], ) diff --git a/skycfg.go b/skycfg.go index 40ffe46..c04e0f6 100644 --- a/skycfg.go +++ b/skycfg.go @@ -358,8 +358,9 @@ type ExecOption interface { } type execOptions struct { - vars *starlark.Dict - funcName string + vars *starlark.Dict + funcName string + flattenLists bool } type fnExecOption func(*execOptions) @@ -382,6 +383,13 @@ func WithEntryPoint(name string) ExecOption { }) } +// WithFlattenLists flatten lists one layer deep ([1, [2,3]] -> [1, 2, 3]) +func WithFlattenLists() ExecOption { + return fnExecOption(func(opts *execOptions) { + opts.flattenLists = true + }) +} + // Main executes main() or a custom entry point function from the top-level Skycfg config // module, which is expected to return either None or a list of Protobuf messages. func (c *Config) Main(ctx context.Context, opts ...ExecOption) ([]proto.Message, error) { @@ -426,11 +434,24 @@ func (c *Config) Main(ctx context.Context, opts ...ExecOption) ([]proto.Message, var msgs []proto.Message for ii := 0; ii < mainList.Len(); ii++ { maybeMsg := mainList.Index(ii) - msg, ok := AsProtoMessage(maybeMsg) - if !ok { - return nil, fmt.Errorf("%q returned something that's not a protobuf (a %s)", parsedOpts.funcName, maybeMsg.Type()) + // Only flatten but not flatten deep. This will flatten out, in order, lists within main list and append the + // message into msgs + if maybeMsgList, ok := maybeMsg.(*starlark.List); parsedOpts.flattenLists && ok { + for iii := 0; iii < maybeMsgList.Len(); iii++ { + maybeNestedMsg := maybeMsgList.Index(iii) + msg, ok := AsProtoMessage(maybeNestedMsg) + if !ok { + return nil, fmt.Errorf("%q returned something that's not a protobuf (a %s) within a nested list", parsedOpts.funcName, maybeNestedMsg.Type()) + } + msgs = append(msgs, msg) + } + } else { + msg, ok := AsProtoMessage(maybeMsg) + if !ok { + return nil, fmt.Errorf("%q returned something that's not a protobuf (a %s)", parsedOpts.funcName, maybeMsg.Type()) + } + msgs = append(msgs, msg) } - msgs = append(msgs, msg) } return msgs, nil } diff --git a/skycfg_test.go b/skycfg_test.go index 8adbbba..3fff5cd 100644 --- a/skycfg_test.go +++ b/skycfg_test.go @@ -176,6 +176,26 @@ def not_main(ctx): msg.f_string = "12345" return [msg] +`, + "test13.sky": ` +test_proto = proto.package("skycfg.test_proto") + +# nested list +def main(ctx): + msg = test_proto.MessageV2() + msg.f_int64 = 12345 + msg.f_string = "12345" + + msg2 = test_proto.MessageV2() + msg2.f_int64 = 123456 + msg2.f_string = "123456" + + msg3 = test_proto.MessageV2() + msg3.f_int64 = 1234567 + msg3.f_string = "1234567" + + innerlist = [msg, msg2] + return [msg3, innerlist] `, } @@ -194,12 +214,13 @@ func (loader *testLoader) ReadFile(ctx context.Context, path string) ([]byte, er } type endToEndTestCase struct { - caseName string - fileToLoad string - vars starlark.StringDict - expLoadErr bool - expExecErr bool - expProtos []proto.Message + caseName string + fileToLoad string + vars starlark.StringDict + expLoadErr bool + expExecErr bool + expProtos []proto.Message + execOptions []skycfg.ExecOption } type ExecSkycfg func(config *skycfg.Config, testCase endToEndTestCase) ([]proto.Message, error) @@ -326,6 +347,25 @@ func TestSkycfgEndToEnd(t *testing.T) { }, }, }, + endToEndTestCase{ + caseName: "flatten nested list", + fileToLoad: "test13.sky", + expProtos: []proto.Message{ + &pb.MessageV2{ + FInt64: proto.Int64(1234567), + FString: proto.String("1234567"), + }, + &pb.MessageV2{ + FInt64: proto.Int64(12345), + FString: proto.String("12345"), + }, + &pb.MessageV2{ + FInt64: proto.Int64(123456), + FString: proto.String("123456"), + }, + }, + execOptions: []skycfg.ExecOption{skycfg.WithFlattenLists()}, + }, endToEndTestCase{ caseName: "value err when attempting to autobox a too large integer into Int32Value", fileToLoad: "test8.sky", @@ -349,7 +389,8 @@ func TestSkycfgEndToEnd(t *testing.T) { } fnExecSkycfg := ExecSkycfg(func(config *skycfg.Config, testCase endToEndTestCase) ([]proto.Message, error) { - return config.Main(context.Background(), skycfg.WithVars(testCase.vars)) + execOptions := append(testCase.execOptions, skycfg.WithVars(testCase.vars)) + return config.Main(context.Background(), execOptions...) }) runTestCases(t, testCases, fnExecSkycfg) }