Protobuf API V2: Add map, repeated, type conversions

go/protomodule/BUILD

@@ -5,8 +5,11 @@ go_library(
     srcs = [
         "protomodule.go",
         "protomodule_enum.go",
+        "protomodule_list.go",
+        "protomodule_map.go",
         "protomodule_message_type.go",
         "protomodule_package.go",
+        "type_conversions.go",
     ],
     importpath = "github.com/stripe/skycfg/go/protomodule",
     visibility = ["//visibility:public"],

go/protomodule/protomodule_enum.go

@@ -93,6 +93,10 @@ func (v *protoEnumValue) Hash() (uint32, error) {
 	return starlark.MakeInt64(int64(v.value.Number())).Hash()
 }
 
+func (v *protoEnumValue) enumNumber() protoreflect.EnumNumber {
+	return v.value.Number()
+}
+
 func (v *protoEnumValue) CompareSameType(op syntax.Token, y starlark.Value, depth int) (bool, error) {
 	other := y.(*protoEnumValue)
 	switch op {

go/protomodule/protomodule_list.go

@@ -0,0 +1,191 @@
+// Copyright 2020 The Skycfg 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.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package protomodule
+
+import (
+	"fmt"
+
+	"go.starlark.net/starlark"
+	"go.starlark.net/syntax"
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+var allowedListMethods = map[string]func(*protoRepeated) starlark.Value{
+	"clear":  nil,
+	"append": (*protoRepeated).wrapAppend,
+	"extend": (*protoRepeated).wrapExtend,
+}
+
+// protoRepeated wraps an underlying starlark.List to provide typechecking on
+// wrties
+//
+// starlark.List is heterogeneous, where protoRepeated enforces all values
+// conform to the given fieldDesc
+type protoRepeated struct {
+	fieldDesc protoreflect.FieldDescriptor
+	list      *starlark.List
+}
+
+var _ starlark.Value = (*protoRepeated)(nil)
+var _ starlark.Iterable = (*protoRepeated)(nil)
+var _ starlark.Sequence = (*protoRepeated)(nil)
+var _ starlark.Indexable = (*protoRepeated)(nil)
+var _ starlark.HasAttrs = (*protoRepeated)(nil)
+var _ starlark.HasSetIndex = (*protoRepeated)(nil)
+var _ starlark.HasBinary = (*protoRepeated)(nil)
+var _ starlark.Comparable = (*protoRepeated)(nil)
+
+func newProtoRepeated(fieldDesc protoreflect.FieldDescriptor) *protoRepeated {
+	return &protoRepeated{fieldDesc, starlark.NewList(nil)}
+}
+
+func newProtoRepeatedFromList(fieldDesc protoreflect.FieldDescriptor, l *starlark.List) (*protoRepeated, error) {
+	out := &protoRepeated{fieldDesc, l}
+	for i := 0; i < l.Len(); i++ {
+		err := scalarTypeCheck(fieldDesc, l.Index(i))
+		if err != nil {
+			return nil, err
+		}
+	}
+	return out, nil
+}
+
+func (r *protoRepeated) Attr(name string) (starlark.Value, error) {
+	wrapper, ok := allowedListMethods[name]
+	if !ok {
+		return nil, nil
+	}
+	if wrapper != nil {
+		return wrapper(r), nil
+	}
+	return r.list.Attr(name)
+}
+
+func (r *protoRepeated) AttrNames() []string                 { return r.list.AttrNames() }
+func (r *protoRepeated) Freeze()                             { r.list.Freeze() }
+func (r *protoRepeated) Hash() (uint32, error)               { return r.list.Hash() }
+func (r *protoRepeated) Index(i int) starlark.Value          { return r.list.Index(i) }
+func (r *protoRepeated) Iterate() starlark.Iterator          { return r.list.Iterate() }
+func (r *protoRepeated) Len() int                            { return r.list.Len() }
+func (r *protoRepeated) Slice(x, y, step int) starlark.Value { return r.list.Slice(x, y, step) }
+func (r *protoRepeated) String() string                      { return r.list.String() }
+func (r *protoRepeated) Truth() starlark.Bool                { return r.list.Truth() }
+
+func (r *protoRepeated) Type() string {
+	return fmt.Sprintf("list<%s>", typeName(r.fieldDesc))
+}
+
+func (r *protoRepeated) CompareSameType(op syntax.Token, y starlark.Value, depth int) (bool, error) {
+	other, ok := y.(*protoRepeated)
+	if !ok {
+		return false, nil
+	}
+
+	return starlark.CompareDepth(op, r.list, other.list, depth)
+}
+
+func (r *protoRepeated) Append(v starlark.Value) error {
+	err := scalarTypeCheck(r.fieldDesc, v)
+	if err != nil {
+		return err
+	}
+
+	err = r.list.Append(v)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (r *protoRepeated) SetIndex(i int, v starlark.Value) error {
+	err := scalarTypeCheck(r.fieldDesc, v)
+	if err != nil {
+		return err
+	}
+
+	err = r.list.SetIndex(i, v)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (r *protoRepeated) Extend(iterable starlark.Iterable) error {
+	iter := iterable.Iterate()
+	defer iter.Done()
+
+	var val starlark.Value
+	for iter.Next(&val) {
+		err := r.Append(val)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (r *protoRepeated) Binary(op syntax.Token, y starlark.Value, side starlark.Side) (starlark.Value, error) {
+	if op == syntax.PLUS {
+		if side == starlark.Left {
+			switch y := y.(type) {
+			case *starlark.List:
+				return starlark.Binary(op, r.list, y)
+			case *protoRepeated:
+				return starlark.Binary(op, r.list, y.list)
+			}
+			return nil, nil
+		}
+		if side == starlark.Right {
+			if _, ok := y.(*starlark.List); ok {
+				return starlark.Binary(op, y, r.list)
+			}
+			return nil, nil
+		}
+	}
+	return nil, nil
+}
+
+func (r *protoRepeated) wrapAppend() starlark.Value {
+	impl := func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+		var val starlark.Value
+		if err := starlark.UnpackPositionalArgs("append", args, kwargs, 1, &val); err != nil {
+			return nil, err
+		}
+		if err := r.Append(val); err != nil {
+			return nil, err
+		}
+		return starlark.None, nil
+	}
+	return starlark.NewBuiltin("append", impl).BindReceiver(r)
+}
+
+func (r *protoRepeated) wrapExtend() starlark.Value {
+	impl := func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+		var val starlark.Iterable
+		if err := starlark.UnpackPositionalArgs("extend", args, kwargs, 1, &val); err != nil {
+			return nil, err
+		}
+		if err := r.Extend(val); err != nil {
+			return nil, err
+		}
+		return starlark.None, nil
+	}
+	return starlark.NewBuiltin("extend", impl).BindReceiver(r)
+}

go/protomodule/protomodule_map.go

@@ -0,0 +1,173 @@
+// Copyright 2020 The Skycfg 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.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package protomodule
+
+import (
+	"fmt"
+
+	"go.starlark.net/starlark"
+	"go.starlark.net/syntax"
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+var allowedDictMethods = map[string]func(*protoMap) starlark.Value{
+	"clear":      nil,
+	"get":        nil,
+	"items":      nil,
+	"keys":       nil,
+	"setdefault": (*protoMap).wrapSetDefault,
+	"update":     (*protoMap).wrapUpdate,
+	"values":     nil,
+}
+
+// protoMap wraps an underlying starlark.Dict to enforce typechecking
+type protoMap struct {
+	mapKey   protoreflect.FieldDescriptor
+	mapValue protoreflect.FieldDescriptor
+	dict     *starlark.Dict
+}
+
+var _ starlark.Value = (*protoMap)(nil)
+var _ starlark.Iterable = (*protoMap)(nil)
+var _ starlark.Sequence = (*protoMap)(nil)
+var _ starlark.HasAttrs = (*protoMap)(nil)
+var _ starlark.HasSetKey = (*protoMap)(nil)
+var _ starlark.Comparable = (*protoMap)(nil)
+
+func newProtoMap(mapKey protoreflect.FieldDescriptor, mapValue protoreflect.FieldDescriptor) *protoMap {
+	return &protoMap{
+		mapKey:   mapKey,
+		mapValue: mapValue,
+		dict:     starlark.NewDict(0),
+	}
+}
+
+func newProtoMapFromDict(mapKey protoreflect.FieldDescriptor, mapValue protoreflect.FieldDescriptor, d *starlark.Dict) (*protoMap, error) {
+	out := &protoMap{
+		mapKey:   mapKey,
+		mapValue: mapValue,
+		dict:     d,
+	}
+
+	for _, item := range d.Items() {
+		err := out.typeCheck(item[0], item[1])
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return out, nil
+}
+
+func (m *protoMap) Attr(name string) (starlark.Value, error) {
+	wrapper, ok := allowedDictMethods[name]
+	if !ok {
+		return nil, nil
+	}
+	if wrapper != nil {
+		return wrapper(m), nil
+	}
+	return m.dict.Attr(name)
+}
+
+func (m *protoMap) AttrNames() []string                                { return m.dict.AttrNames() }
+func (m *protoMap) Freeze()                                            { m.dict.Freeze() }
+func (m *protoMap) Hash() (uint32, error)                              { return m.dict.Hash() }
+func (m *protoMap) Get(k starlark.Value) (starlark.Value, bool, error) { return m.dict.Get(k) }
+func (m *protoMap) Iterate() starlark.Iterator                         { return m.dict.Iterate() }
+func (m *protoMap) Len() int                                           { return m.dict.Len() }
+func (m *protoMap) String() string                                     { return m.dict.String() }
+func (m *protoMap) Truth() starlark.Bool                               { return m.dict.Truth() }
+func (m *protoMap) Items() []starlark.Tuple                            { return m.dict.Items() }
+
+func (m *protoMap) Type() string {
+	return fmt.Sprintf("map<%s, %s>", typeName(m.mapKey), typeName(m.mapValue))
+}
+
+func (m *protoMap) CompareSameType(op syntax.Token, y starlark.Value, depth int) (bool, error) {
+	other, ok := y.(*protoMap)
+	if !ok {
+		return false, nil
+	}
+
+	return starlark.CompareDepth(op, m.dict, other.dict, depth)
+}
+
+func (m *protoMap) wrapSetDefault() starlark.Value {
+	impl := func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+		var key, defaultValue starlark.Value = nil, starlark.None
+		if err := starlark.UnpackPositionalArgs("setdefault", args, kwargs, 1, &key, &defaultValue); err != nil {
+			return nil, err
+		}
+		if val, ok, err := m.dict.Get(key); err != nil {
+			return nil, err
+		} else if ok {
+			return val, nil
+		}
+		return defaultValue, m.SetKey(key, defaultValue)
+	}
+	return starlark.NewBuiltin("setdefault", impl).BindReceiver(m)
+}
+
+func (m *protoMap) wrapUpdate() starlark.Value {
+	impl := func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+		// Use the underlying starlark `dict.update()` to get a Dict containing
+		// all the new values, so we don't have to recreate the API here. After
+		// the temp dict is constructed, type check.
+		tempDict := &starlark.Dict{}
+		tempUpdate, _ := tempDict.Attr("update")
+		if _, err := starlark.Call(thread, tempUpdate, args, kwargs); err != nil {
+			return nil, err
+		}
+		for _, item := range tempDict.Items() {
+			if err := m.SetKey(item[0], item[1]); err != nil {
+				return nil, err
+			}
+		}
+
+		return starlark.None, nil
+	}
+	return starlark.NewBuiltin("update", impl).BindReceiver(m)
+}
+
+func (m *protoMap) SetKey(k, v starlark.Value) error {
+	err := m.typeCheck(k, v)
+	if err != nil {
+		return err
+	}
+
+	err = m.dict.SetKey(k, v)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (m *protoMap) typeCheck(k, v starlark.Value) error {
+	err := scalarTypeCheck(m.mapKey, k)
+	if err != nil {
+		return err
+	}
+
+	err = scalarTypeCheck(m.mapValue, v)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}