1// The star2proto command executes a Starlark file and prints a protocol 2// message, which it expects to find in a module-level variable named 'result'. 3// 4// THIS COMMAND IS EXPERIMENTAL AND ITS INTERFACE MAY CHANGE. 5package main 6 7// TODO(adonovan): add features to make this a useful tool for querying, 8// converting, and building messages in proto, JSON, and YAML. 9// - define operations for reading and writing files. 10// - support (e.g.) querying a proto file given a '-e expr' flag. 11// This will need a convenient way to put the relevant descriptors in scope. 12 13import ( 14 "flag" 15 "fmt" 16 "io/ioutil" 17 "log" 18 "os" 19 "strings" 20 21 starlarkproto "go.starlark.net/lib/proto" 22 "go.starlark.net/resolve" 23 "go.starlark.net/starlark" 24 "go.starlark.net/starlarkjson" 25 "google.golang.org/protobuf/encoding/protojson" 26 "google.golang.org/protobuf/encoding/prototext" 27 "google.golang.org/protobuf/proto" 28 "google.golang.org/protobuf/reflect/protodesc" 29 "google.golang.org/protobuf/reflect/protoreflect" 30 "google.golang.org/protobuf/reflect/protoregistry" 31 "google.golang.org/protobuf/types/descriptorpb" 32) 33 34// flags 35var ( 36 outputFlag = flag.String("output", "text", "output format (text, wire, json)") 37 varFlag = flag.String("var", "result", "the variable to output") 38 descriptors = flag.String("descriptors", "", "comma-separated list of names of files containing proto.FileDescriptorProto messages") 39) 40 41// Starlark dialect flags 42func init() { 43 flag.BoolVar(&resolve.AllowFloat, "fp", true, "allow floating-point numbers") 44 flag.BoolVar(&resolve.AllowSet, "set", resolve.AllowSet, "allow set data type") 45 flag.BoolVar(&resolve.AllowLambda, "lambda", resolve.AllowLambda, "allow lambda expressions") 46 flag.BoolVar(&resolve.AllowNestedDef, "nesteddef", resolve.AllowNestedDef, "allow nested def statements") 47} 48 49func main() { 50 log.SetPrefix("star2proto: ") 51 log.SetFlags(0) 52 flag.Parse() 53 if len(flag.Args()) != 1 { 54 fatalf("requires a single Starlark file name") 55 } 56 filename := flag.Args()[0] 57 58 // By default, use the linked-in descriptors 59 // (very few in star2proto, e.g. descriptorpb itself). 60 pool := protoregistry.GlobalFiles 61 62 // Load a user-provided FileDescriptorSet produced by a command such as: 63 // $ protoc --descriptor_set_out=foo.fds foo.proto 64 if *descriptors != "" { 65 var fdset descriptorpb.FileDescriptorSet 66 for i, filename := range strings.Split(*descriptors, ",") { 67 data, err := ioutil.ReadFile(filename) 68 if err != nil { 69 log.Fatalf("--descriptors[%d]: %s", i, err) 70 } 71 // Accumulate into the repeated field of FileDescriptors. 72 if err := (proto.UnmarshalOptions{Merge: true}).Unmarshal(data, &fdset); err != nil { 73 log.Fatalf("%s does not contain a proto2.FileDescriptorSet: %v", filename, err) 74 } 75 } 76 77 files, err := protodesc.NewFiles(&fdset) 78 if err != nil { 79 log.Fatalf("protodesc.NewFiles: could not build FileDescriptor index: %v", err) 80 } 81 pool = files 82 } 83 84 // Execute the Starlark file. 85 thread := &starlark.Thread{ 86 Print: func(_ *starlark.Thread, msg string) { fmt.Println(msg) }, 87 } 88 starlarkproto.SetPool(thread, pool) 89 predeclared := starlark.StringDict{ 90 "proto": starlarkproto.Module, 91 "json": starlarkjson.Module, 92 } 93 globals, err := starlark.ExecFile(thread, filename, nil, predeclared) 94 if err != nil { 95 if evalErr, ok := err.(*starlark.EvalError); ok { 96 fatalf("%s", evalErr.Backtrace()) 97 } else { 98 fatalf("%s", err) 99 } 100 } 101 102 // Print the output variable as a message. 103 // TODO(adonovan): this is clumsy. 104 // Let the user call print(), or provide an expression on the command line. 105 result, ok := globals[*varFlag] 106 if !ok { 107 fatalf("%s must define a module-level variable named %q", filename, *varFlag) 108 } 109 msgwrap, ok := result.(*starlarkproto.Message) 110 if !ok { 111 fatalf("got %s, want proto.Message, for %q", result.Type(), *varFlag) 112 } 113 msg := msgwrap.Message() 114 115 // -output 116 var marshal func(protoreflect.ProtoMessage) ([]byte, error) 117 switch *outputFlag { 118 case "wire": 119 marshal = proto.Marshal 120 121 case "text": 122 marshal = prototext.MarshalOptions{Multiline: true, Indent: "\t"}.Marshal 123 124 case "json": 125 marshal = protojson.MarshalOptions{Multiline: true, Indent: "\t"}.Marshal 126 127 default: 128 fatalf("unsupported -output format: %s", *outputFlag) 129 } 130 data, err := marshal(msg) 131 if err != nil { 132 fatalf("%s", err) 133 } 134 os.Stdout.Write(data) 135} 136 137func fatalf(format string, args ...interface{}) { 138 fmt.Fprintf(os.Stderr, "star2proto: ") 139 fmt.Fprintf(os.Stderr, format, args...) 140 fmt.Fprintln(os.Stderr) 141 os.Exit(1) 142} 143