Add flatten on eval (#104)

* add flatten * update protobuf for zlib + comment * fix test * add flatten as an option * remove conditional
stripe · Feb 3, 2023 · 35ad986 · 35ad986
1 parent 4571141
commit 35ad986
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 17 deletions.
diff --git 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
@@ -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
@@ -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)
 }