diff options
Diffstat (limited to 'lib/proto/cmd/star2proto/star2proto.go')
-rw-r--r-- | lib/proto/cmd/star2proto/star2proto.go | 142 |
1 files changed, 142 insertions, 0 deletions
diff --git a/lib/proto/cmd/star2proto/star2proto.go b/lib/proto/cmd/star2proto/star2proto.go new file mode 100644 index 0000000..7911723 --- /dev/null +++ b/lib/proto/cmd/star2proto/star2proto.go @@ -0,0 +1,142 @@ +// The star2proto command executes a Starlark file and prints a protocol +// message, which it expects to find in a module-level variable named 'result'. +// +// THIS COMMAND IS EXPERIMENTAL AND ITS INTERFACE MAY CHANGE. +package main + +// TODO(adonovan): add features to make this a useful tool for querying, +// converting, and building messages in proto, JSON, and YAML. +// - define operations for reading and writing files. +// - support (e.g.) querying a proto file given a '-e expr' flag. +// This will need a convenient way to put the relevant descriptors in scope. + +import ( + "flag" + "fmt" + "io/ioutil" + "log" + "os" + "strings" + + starlarkproto "go.starlark.net/lib/proto" + "go.starlark.net/resolve" + "go.starlark.net/starlark" + "go.starlark.net/starlarkjson" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/encoding/prototext" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protodesc" + "google.golang.org/protobuf/reflect/protoreflect" + "google.golang.org/protobuf/reflect/protoregistry" + "google.golang.org/protobuf/types/descriptorpb" +) + +// flags +var ( + outputFlag = flag.String("output", "text", "output format (text, wire, json)") + varFlag = flag.String("var", "result", "the variable to output") + descriptors = flag.String("descriptors", "", "comma-separated list of names of files containing proto.FileDescriptorProto messages") +) + +// Starlark dialect flags +func init() { + flag.BoolVar(&resolve.AllowFloat, "fp", true, "allow floating-point numbers") + flag.BoolVar(&resolve.AllowSet, "set", resolve.AllowSet, "allow set data type") + flag.BoolVar(&resolve.AllowLambda, "lambda", resolve.AllowLambda, "allow lambda expressions") + flag.BoolVar(&resolve.AllowNestedDef, "nesteddef", resolve.AllowNestedDef, "allow nested def statements") +} + +func main() { + log.SetPrefix("star2proto: ") + log.SetFlags(0) + flag.Parse() + if len(flag.Args()) != 1 { + fatalf("requires a single Starlark file name") + } + filename := flag.Args()[0] + + // By default, use the linked-in descriptors + // (very few in star2proto, e.g. descriptorpb itself). + pool := protoregistry.GlobalFiles + + // Load a user-provided FileDescriptorSet produced by a command such as: + // $ protoc --descriptor_set_out=foo.fds foo.proto + if *descriptors != "" { + var fdset descriptorpb.FileDescriptorSet + for i, filename := range strings.Split(*descriptors, ",") { + data, err := ioutil.ReadFile(filename) + if err != nil { + log.Fatalf("--descriptors[%d]: %s", i, err) + } + // Accumulate into the repeated field of FileDescriptors. + if err := (proto.UnmarshalOptions{Merge: true}).Unmarshal(data, &fdset); err != nil { + log.Fatalf("%s does not contain a proto2.FileDescriptorSet: %v", filename, err) + } + } + + files, err := protodesc.NewFiles(&fdset) + if err != nil { + log.Fatalf("protodesc.NewFiles: could not build FileDescriptor index: %v", err) + } + pool = files + } + + // Execute the Starlark file. + thread := &starlark.Thread{ + Print: func(_ *starlark.Thread, msg string) { fmt.Println(msg) }, + } + starlarkproto.SetPool(thread, pool) + predeclared := starlark.StringDict{ + "proto": starlarkproto.Module, + "json": starlarkjson.Module, + } + globals, err := starlark.ExecFile(thread, filename, nil, predeclared) + if err != nil { + if evalErr, ok := err.(*starlark.EvalError); ok { + fatalf("%s", evalErr.Backtrace()) + } else { + fatalf("%s", err) + } + } + + // Print the output variable as a message. + // TODO(adonovan): this is clumsy. + // Let the user call print(), or provide an expression on the command line. + result, ok := globals[*varFlag] + if !ok { + fatalf("%s must define a module-level variable named %q", filename, *varFlag) + } + msgwrap, ok := result.(*starlarkproto.Message) + if !ok { + fatalf("got %s, want proto.Message, for %q", result.Type(), *varFlag) + } + msg := msgwrap.Message() + + // -output + var marshal func(protoreflect.ProtoMessage) ([]byte, error) + switch *outputFlag { + case "wire": + marshal = proto.Marshal + + case "text": + marshal = prototext.MarshalOptions{Multiline: true, Indent: "\t"}.Marshal + + case "json": + marshal = protojson.MarshalOptions{Multiline: true, Indent: "\t"}.Marshal + + default: + fatalf("unsupported -output format: %s", *outputFlag) + } + data, err := marshal(msg) + if err != nil { + fatalf("%s", err) + } + os.Stdout.Write(data) +} + +func fatalf(format string, args ...interface{}) { + fmt.Fprintf(os.Stderr, "star2proto: ") + fmt.Fprintf(os.Stderr, format, args...) + fmt.Fprintln(os.Stderr) + os.Exit(1) +} |