Added support custom type and fixed bugs

go.mod

@@ -19,7 +19,8 @@ require (
 	github.com/dustmop/soup v1.1.2-0.20190516214245-38228baa104e
 	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
 	github.com/paulmach/orb v0.1.5
-	github.com/stretchr/testify v1.5.1 // indirect
+	github.com/pkg/errors v0.8.1
+	github.com/stretchr/testify v1.5.1
 	go.starlark.net v0.0.0-20200330013621-be5394c419b6
 	golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e
 	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f // indirect

go.sum

@@ -24,6 +24,7 @@ github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWb
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/paulmach/orb v0.1.5 h1:GUcATabvxciqEzGd+c01/9ek3B6pUp9OdcIHFSDDSSg=
 github.com/paulmach/orb v0.1.5/go.mod h1:pPwxxs3zoAyosNSbNKn1jiXV2+oovRDObDKfTvRegDI=
+github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=

util/util.go

@@ -4,7 +4,9 @@ import (
 	"fmt"
 	"strconv"
 
+	"github.com/pkg/errors"
 	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
 )
 
 // asString unquotes a starlark string value
@@ -24,7 +26,7 @@ func IsEmptyString(s starlark.String) bool {
 func Unmarshal(x starlark.Value) (val interface{}, err error) {
 	switch v := x.(type) {
 	case starlark.NoneType:
-		val = v
+		val = nil
 	case starlark.Bool:
 		val = v.Truth() == starlark.True
 	case starlark.Int:
@@ -41,7 +43,11 @@ func Unmarshal(x starlark.Value) (val interface{}, err error) {
 		var (
 			dictVal starlark.Value
 			pval    interface{}
-			value   = map[string]interface{}{}
+			kval    interface{}
+			keys    []interface{}
+			vals    []interface{}
+			// key as interface if found one key is not a string
+			ki bool
 		)
 
 		for _, k := range v.Keys() {
@@ -52,18 +58,44 @@ func Unmarshal(x starlark.Value) (val interface{}, err error) {
 
 			pval, err = Unmarshal(dictVal)
 			if err != nil {
+				err = fmt.Errorf("unmarshaling starlark value: %w", err)
 				return
 			}
 
-			var str string
-			str, err = asString(k)
+			kval, err = Unmarshal(k)
 			if err != nil {
+				err = fmt.Errorf("unmarshaling starlark key: %w", err)
 				return
 			}
 
-			value[str] = pval
+			if _, ok := kval.(string); !ok {
+				// found key as not a string
+				ki = true
+			}
+
+			keys = append(keys, kval)
+			vals = append(vals, pval)
+		}
+
+		// prepare result
+
+		rs := map[string]interface{}{}
+		ri := map[interface{}]interface{}{}
+
+		for i, key := range keys {
+			// key as interface
+			if ki {
+				ri[key] = vals[i]
+			} else {
+				rs[key.(string)] = vals[i]
+			}
+		}
+
+		if ki {
+			val = ri // map[interface{}]interface{}
+		} else {
+			val = rs // map[string]interface{}
 		}
-		val = value
 	case *starlark.List:
 		var (
 			i       int
@@ -101,6 +133,17 @@ func Unmarshal(x starlark.Value) (val interface{}, err error) {
 	case *starlark.Set:
 		fmt.Println("errnotdone: SET")
 		err = fmt.Errorf("sets aren't yet supported")
+	case *starlarkstruct.Struct:
+		if _var, ok := v.Constructor().(Unmarshaler); ok {
+			err = _var.UnmarshalStarlark(x)
+			if err != nil {
+				err = errors.Wrapf(err, "failed marshal %q to Starlark object", v.Constructor().Type())
+				return
+			}
+			val = _var
+		} else {
+			err = fmt.Errorf("constructor object from *starlarkstruct.Struct not supported Marshaler to starlark object: %s", v.Constructor().Type())
+		}
 	default:
 		fmt.Println("errbadtype:", x.Type())
 		err = fmt.Errorf("unrecognized starlark type: %s", x.Type())
@@ -182,8 +225,22 @@ func Marshal(data interface{}) (v starlark.Value, err error) {
 			}
 		}
 		v = dict
+	case Marshaler:
+		v, err = x.MarshalStarlark()
 	default:
 		return starlark.None, fmt.Errorf("unrecognized type: %#v", x)
 	}
 	return
 }
+
+// Unmarshaler is the interface use to unmarshal starlark custom types.
+type Unmarshaler interface {
+	// UnmarshalStarlark unmarshal a starlark object to custom type.
+	UnmarshalStarlark(starlark.Value) error
+}
+
+// Marshaler is the interface use to marshal starlark custom types.
+type Marshaler interface {
+	// MarshalStarlark marshal a custom type to starlark object.
+	MarshalStarlark() (starlark.Value, error)
+}

util/util_test.go

@@ -1,9 +1,12 @@
 package util
 
 import (
+	"fmt"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
 )
 
 func TestIsEmptyString(t *testing.T) {
@@ -46,10 +49,15 @@ func TestMarshal(t *testing.T) {
 	expectedIntDict := starlark.NewDict(1)
 	expectedIntDict.SetKey(starlark.MakeInt(42*2), starlark.MakeInt(42))
 
+	ct, _ := (&customType{42}).MarshalStarlark()
+	expectedStrDictCustomType := starlark.NewDict(2)
+	expectedStrDictCustomType.SetKey(starlark.String("foo"), starlark.MakeInt(42))
+	expectedStrDictCustomType.SetKey(starlark.String("bar"), ct)
+
 	cases := []struct {
-		in  interface{}
-		got starlark.Value
-		err string
+		in   interface{}
+		want starlark.Value
+		err  string
 	}{
 		{nil, starlark.None, ""},
 		{true, starlark.True, ""},
@@ -70,33 +78,129 @@ func TestMarshal(t *testing.T) {
 		{map[string]interface{}{"foo": 42}, expectedStringDict, ""},
 		{map[interface{}]interface{}{"foo": 42}, expectedStringDict, ""},
 		{map[interface{}]interface{}{42 * 2: 42}, expectedIntDict, ""},
+		{&customType{42}, ct, ""},
+		{map[string]interface{}{"foo": 42, "bar": &customType{42}}, expectedStrDictCustomType, ""},
+		{map[interface{}]interface{}{"foo": 42, "bar": &customType{42}}, expectedStrDictCustomType, ""},
+		{[]interface{}{42, &customType{42}}, starlark.NewList([]starlark.Value{starlark.MakeInt(42), ct}), ""},
+		{&invalidCustomType{42}, starlark.None, "unrecognized type: &util.invalidCustomType{Foo:42}"},
 	}
 
 	for i, c := range cases {
 		got, err := Marshal(c.in)
 		if !(err == nil && c.err == "" || err != nil && err.Error() == c.err) {
-			t.Errorf("case %d error mismatch. expected: '%s', got: '%s'", i, c.err, err)
+			t.Errorf("case %d error mismatch. expected: %q, got: %q (%T -> %T)", i, c.err, err, c.in, c.want)
 			continue
 		}
 
-		if list, ok := c.got.(*starlark.List); ok {
-			if list.String() != got.String() {
-				t.Errorf("case %d. expected: '%s', got: '%s'", i, c.got, got)
-			}
+		assert.EqualValues(t, c.want, got, "case %d: %T -> %T", i, c.in, c.want)
+	}
+}
 
-			continue
-		}
+func TestUnmarshal(t *testing.T) {
+	strDict := starlark.NewDict(1)
+	strDict.SetKey(starlark.String("foo"), starlark.MakeInt(42))
 
-		if dict, ok := c.got.(*starlark.Dict); ok {
-			if dict.String() != got.String() {
-				t.Errorf("case %d. expected: '%s', got: '%s'", i, c.got, got)
-			}
+	intDict := starlark.NewDict(1)
+	intDict.SetKey(starlark.MakeInt(42*2), starlark.MakeInt(42))
 
+	ct, _ := (&customType{42}).MarshalStarlark()
+	strDictCT := starlark.NewDict(2)
+	strDictCT.SetKey(starlark.String("foo"), starlark.MakeInt(42))
+	strDictCT.SetKey(starlark.String("bar"), ct)
+
+	cases := []struct {
+		in   starlark.Value
+		want interface{}
+		err  string
+	}{
+		{starlark.None, nil, ""},
+		{starlark.True, true, ""},
+		{starlark.String("foo"), "foo", ""},
+		{starlark.MakeInt(42), 42, ""},
+		{starlark.MakeInt(42), int8(42), ""},
+		{starlark.MakeInt(42), int16(42), ""},
+		{starlark.MakeInt(42), int32(42), ""},
+		{starlark.MakeInt(42), int64(42), ""},
+		{starlark.MakeUint(42), uint(42), ""},
+		{starlark.MakeUint(42), uint8(42), ""},
+		{starlark.MakeUint(42), uint16(42), ""},
+		{starlark.MakeUint(42), uint32(42), ""},
+		{starlark.MakeUint64(42), uint64(42), ""},
+		{starlark.Float(42), float32(42), ""},
+		{starlark.Float(42), 42., ""},
+		{starlark.NewList([]starlark.Value{starlark.MakeInt(42)}), []interface{}{42}, ""},
+		{strDict, map[string]interface{}{"foo": 42}, ""},
+		{intDict, map[interface{}]interface{}{42 * 2: 42}, ""},
+		{ct, &customType{42}, ""},
+		{strDictCT, map[string]interface{}{"foo": 42, "bar": &customType{42}}, ""},
+		{starlark.NewList([]starlark.Value{starlark.MakeInt(42), ct}), []interface{}{42, &customType{42}}, ""},
+	}
+
+	for i, c := range cases {
+		got, err := Unmarshal(c.in)
+		if !(err == nil && c.err == "" || err != nil && err.Error() == c.err) {
+			t.Errorf("case %d error mismatch. expected: %q, got: %q %T -> %T", i, c.err, err, c.in, c.want)
 			continue
 		}
 
-		if c.got != got {
-			t.Errorf("case %d. expected: '%s', got: '%s'", i, c.got, got)
-		}
+		assert.EqualValues(t, c.want, got, "case %d: %T -> %T", i, c.in, c.want)
 	}
 }
+
+type invalidCustomType struct {
+	Foo int64
+}
+
+type customType invalidCustomType
+
+func (t *customType) UnmarshalStarlark(v starlark.Value) error {
+	// asserts
+	if v.Type() != "struct" {
+		return fmt.Errorf("not expected top level type, want struct, got %q", v.Type())
+	}
+	if _, ok := v.(*starlarkstruct.Struct).Constructor().(*customType); !ok {
+		return fmt.Errorf("not expected construct type got %T, want %T", v.(*starlarkstruct.Struct).Constructor(), t)
+	}
+
+	// TODO: refactoring transform data
+
+	mustInt64 := func(sv starlark.Value) int64 {
+		i, _ := sv.(starlark.Int).Int64()
+		return i
+	}
+
+	data := starlark.StringDict{}
+	v.(*starlarkstruct.Struct).ToStringDict(data)
+
+	*t = customType{
+		Foo: mustInt64(data["foo"]),
+	}
+	return nil
+}
+
+func (t *customType) MarshalStarlark() (starlark.Value, error) {
+	v := starlarkstruct.FromStringDict(&customType{}, starlark.StringDict{
+		"foo": starlark.MakeInt64(t.Foo),
+	})
+	return v, nil
+}
+
+func (c customType) String() string {
+	return "customType"
+}
+
+func (c customType) Type() string { return "test.customType" }
+
+func (customType) Freeze() {}
+
+func (c customType) Truth() starlark.Bool {
+	return starlark.True
+}
+
+func (c customType) Hash() (uint32, error) {
+	return 0, fmt.Errorf("unhashable: %s", c.Type())
+}
+
+var _ Unmarshaler = (*customType)(nil)
+var _ Marshaler = (*customType)(nil)
+var _ starlark.Value = (*customType)(nil)