[starlark] Add genstruct(...) and ctor(...) builtins.
They can be used to create structs "tagged" with a symbol. It's not quite
like a typed struct (each individual object still can have arbitrary fields),
but pretty close:
mystruct = genstruct("mystruct")
s = mystruct(a=1, b=2)
assert.eq(ctor(s), mystruct)
assert.eq(ctor(1), None)
assert.eq(ctor(struct(a=1, b=2)), "struct")
Can be used to add some sort of type safety to Starlark libraries. Note that
using type(strct) is insufficient, since it always returns "struct".
R=nodir@chromium.org, iannucci@chromium.org
BUG=833946
Change-Id: I65abc0594441367a7ecd3c0faa05c06499d61c30
Reviewed-on: https://chromium-review.googlesource.com/c/1370693
Reviewed-by: Robbie Iannucci <iannucci@chromium.org>
Reviewed-by: Nodir Turakulov <nodir@chromium.org>
Commit-Queue: Vadim Shtayura <vadimsh@chromium.org>
diff --git a/starlark/builtins/struct.go b/starlark/builtins/struct.go
index 867a49c..7a57589 100644
--- a/starlark/builtins/struct.go
+++ b/starlark/builtins/struct.go
@@ -15,12 +15,63 @@
package builtins
import (
+ "fmt"
+
"go.starlark.net/starlark"
"go.starlark.net/starlarkstruct"
)
-// Struct is struct(**kwargs) builtin.
-//
-// def struct(**kwargs):
-// """Returns an immutable object with fields set to given values."""
-var Struct = starlark.NewBuiltin("struct", starlarkstruct.Make)
+var (
+ // Struct is struct(**kwargs) builtin.
+ //
+ // def struct(**kwargs):
+ // """Returns an immutable object with fields set to given values."""
+ Struct = starlark.NewBuiltin("struct", starlarkstruct.Make)
+
+ // GenStruct is genstruct(name) builtin.
+ //
+ // def genstruct(name):
+ // """Returns a callable constructor for "branded" struct instances."""
+ GenStruct = starlark.NewBuiltin("genstruct", func(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+ var name string
+ if err := starlark.UnpackArgs("genstruct", args, kwargs, "name", &name); err != nil {
+ return nil, err
+ }
+ return &ctor{name: name}, nil
+ })
+
+ // Ctor is ctor(obj) builtin.
+ //
+ // def ctor(obj):
+ // """Returns a constructor used to construct this struct or None."""
+ Ctor = starlark.NewBuiltin("ctor", func(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+ var obj starlark.Value
+ if err := starlark.UnpackArgs("ctor", args, kwargs, "obj", &obj); err != nil {
+ return nil, err
+ }
+ if st, ok := obj.(*starlarkstruct.Struct); ok {
+ return st.Constructor(), nil
+ }
+ return starlark.None, nil
+ })
+)
+
+// Ctor is a callable that produces starlark structs that have it as a
+// constructor.
+type ctor struct{ name string }
+
+var _ starlark.Callable = (*ctor)(nil)
+
+func (c *ctor) Name() string { return c.name }
+func (c *ctor) String() string { return c.name }
+func (c *ctor) Type() string { return "ctor" }
+func (c *ctor) Freeze() {} // immutable
+func (c *ctor) Truth() starlark.Bool { return starlark.True }
+func (c *ctor) Hash() (uint32, error) { return 0, fmt.Errorf("unhashable: ctor") }
+
+func (c *ctor) CallInternal(_ *starlark.Thread, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+ if len(args) > 0 {
+ return nil, fmt.Errorf("%s: unexpected positional arguments", c)
+ }
+ return starlarkstruct.FromKeywords(c, kwargs), nil
+}
diff --git a/starlark/builtins/struct_test.go b/starlark/builtins/struct_test.go
new file mode 100644
index 0000000..2fd5e12
--- /dev/null
+++ b/starlark/builtins/struct_test.go
@@ -0,0 +1,54 @@
+// Copyright 2018 The LUCI Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package builtins
+
+import (
+ "testing"
+
+ "go.starlark.net/starlark"
+
+ . "github.com/smartystreets/goconvey/convey"
+)
+
+const successTest = `
+def assert(a, b):
+ if a != b:
+ fail("%s (%s) != %s (%s)" % (a, type(a), b, type(b)))
+
+mystruct = genstruct("mystruct")
+s = mystruct(a=1, b=2)
+assert(ctor(s), mystruct)
+assert(ctor(1), None)
+assert(ctor(struct(a=1, b=2)), "struct")
+`
+
+func TestStruct(t *testing.T) {
+ t.Parallel()
+
+ runScript := func(code string) error {
+ _, err := starlark.ExecFile(&starlark.Thread{}, "main", code, starlark.StringDict{
+ "struct": Struct,
+ "genstruct": GenStruct,
+ "ctor": Ctor,
+
+ "fail": Fail, // for 'assert'
+ })
+ return err
+ }
+
+ Convey("Success", t, func() {
+ So(runScript(successTest), ShouldBeNil)
+ })
+}
