Skip to content

Commit

Permalink
Support structpb too
Browse files Browse the repository at this point in the history
  • Loading branch information
chriso committed Jun 25, 2024
1 parent 219b268 commit 2b61003
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 0 deletions.
60 changes: 60 additions & 0 deletions cli/any.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@ import (
"fmt"
"log/slog"
"strconv"
"strings"

pythonv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/python/v1"
"google.golang.org/protobuf/types/known/anypb"
"google.golang.org/protobuf/types/known/durationpb"
"google.golang.org/protobuf/types/known/emptypb"
"google.golang.org/protobuf/types/known/structpb"
"google.golang.org/protobuf/types/known/timestamppb"
"google.golang.org/protobuf/types/known/wrapperspb"
)
Expand Down Expand Up @@ -67,6 +69,15 @@ func anyString(any *anypb.Any) string {
case *durationpb.Duration:
return mm.AsDuration().String()

case *structpb.Struct:
return structpbStructString(mm)

case *structpb.ListValue:
return structpbListString(mm)

case *structpb.Value:
return structpbValueString(mm)

case *pythonv1.Pickled:
s, err := pythonPickleString(mm.PickledValue)
if err != nil {
Expand All @@ -79,6 +90,55 @@ func anyString(any *anypb.Any) string {
}
}

func structpbStructString(s *structpb.Struct) string {
var b strings.Builder
b.WriteByte('{')
i := 0
for name, value := range s.Fields {
if i > 0 {
b.WriteString(", ")
}
b.WriteString(fmt.Sprintf("%q", name))
b.WriteString(": ")
b.WriteString(structpbValueString(value))
i++
}
b.WriteByte('}')
return b.String()
}

func structpbListString(s *structpb.ListValue) string {
var b strings.Builder
b.WriteByte('[')
for i, value := range s.Values {
if i > 0 {
b.WriteString(", ")
}
b.WriteString(structpbValueString(value))
}
b.WriteByte(']')
return b.String()
}

func structpbValueString(s *structpb.Value) string {
switch v := s.Kind.(type) {
case *structpb.Value_StructValue:
return structpbStructString(v.StructValue)
case *structpb.Value_ListValue:
return structpbListString(v.ListValue)
case *structpb.Value_BoolValue:
return strconv.FormatBool(v.BoolValue)
case *structpb.Value_NumberValue:
return fmt.Sprintf("%v", v.NumberValue)
case *structpb.Value_StringValue:
return fmt.Sprintf("%q", v.StringValue)
case *structpb.Value_NullValue:
return "null"
default:
panic("unreachable")
}
}

func unsupportedAny(any *anypb.Any, err error) string {
if err != nil {
slog.Debug("cannot parse input/output value", "error", err)
Expand Down
37 changes: 37 additions & 0 deletions cli/any_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
"google.golang.org/protobuf/types/known/anypb"
"google.golang.org/protobuf/types/known/durationpb"
"google.golang.org/protobuf/types/known/emptypb"
"google.golang.org/protobuf/types/known/structpb"
"google.golang.org/protobuf/types/known/timestamppb"
"google.golang.org/protobuf/types/known/wrapperspb"
)
Expand Down Expand Up @@ -91,6 +92,34 @@ func TestAnyString(t *testing.T) {
input: asAny(&emptypb.Empty{}),
want: "empty()",
},
{
input: asAny(structpb.NewNullValue()),
want: "null",
},
{
input: asAny(structpb.NewBoolValue(false)),
want: "false",
},
{
input: asAny(structpb.NewNumberValue(1111)),
want: "1111",
},
{
input: asAny(structpb.NewNumberValue(3.14)),
want: "3.14",
},
{
input: asAny(structpb.NewStringValue("foobar")),
want: `"foobar"`,
},
{
input: asStructValue([]any{1, true, "abc", nil, map[string]any{}, []any{}}),
want: `[1, true, "abc", null, {}, []]`,
},
{
input: asStructValue(map[string]any{"foo": []any{"bar", "baz"}}),
want: `{"foo": ["bar", "baz"]}`,
},
} {
t.Run(test.want, func(*testing.T) {
got := anyString(test.input)
Expand All @@ -109,6 +138,14 @@ func asAny(m proto.Message) *anypb.Any {
return any
}

func asStructValue(v any) *anypb.Any {
m, err := structpb.NewValue(v)
if err != nil {
panic(err)
}
return asAny(m)
}

func pickled(b []byte) *anypb.Any {
m := &pythonv1.Pickled{PickledValue: b}
mb, err := proto.Marshal(m)
Expand Down

0 comments on commit 2b61003

Please sign in to comment.