File: star2proto.go

package info (click to toggle)
golang-starlark 0.0~git20240725.42030a7-1
  • links: PTS, VCS
  • area: main
  • in suites: experimental, forky, sid, trixie
  • size: 1,296 kB
  • sloc: makefile: 8; sh: 8
file content (143 lines) | stat: -rw-r--r-- 4,489 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
// 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"
	"log"
	"os"
	"strings"

	"go.starlark.net/lib/json"
	starlarkproto "go.starlark.net/lib/proto"
	"go.starlark.net/resolve"
	"go.starlark.net/starlark"
	"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.AllowSet, "set", resolve.AllowSet, "allow set data type")

	// obsolete, no effect:
	flag.BoolVar(&resolve.AllowFloat, "fp", true, "allow floating-point numbers")
	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 := os.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":  json.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)
}