star2proto.go

Documentation: go.starlark.net/lib/proto/cmd/star2proto

     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.
     5  package 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  
    13  import (
    14  	"flag"
    15  	"fmt"
    16  	"io/ioutil"
    17  	"log"
    18  	"os"
    19  	"strings"
    20  
    21  	"go.starlark.net/lib/json"
    22  	starlarkproto "go.starlark.net/lib/proto"
    23  	"go.starlark.net/resolve"
    24  	"go.starlark.net/starlark"
    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
    35  var (
    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
    42  func 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  
    49  func 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":  json.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  
   137  func 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
View as plain text