go.starlark.net@v0.0.0-20231101134539-556fd59b42f6/lib/proto/cmd/star2proto/star2proto.go (about)

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