lib/proto/cmd/star2proto/star2proto.go

// 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)
}