diff --git a/go.mod b/go.mod index f110a51..d7bce87 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index d62bb78..60423e2 100644 --- a/go.sum +++ b/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= diff --git a/util/util.go b/util/util.go index ffe48d0..185b989 100644 --- a/util/util.go +++ b/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) +} diff --git a/util/util_test.go b/util/util_test.go index 78ea566..76b8a06 100644 --- a/util/util_test.go +++ b/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)