forked from apache/beam
-
Notifications
You must be signed in to change notification settings - Fork 0
/
starcgen.go
199 lines (180 loc) · 6.63 KB
/
starcgen.go
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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// starcgen is a tool to generate specialized type assertion shims to be
// used in Apache Beam Go SDK pipelines instead of the default reflection shim.
// This is done through static analysis of go sources for the package in question.
//
// The generated type assertion shims have much better performance than the default
// reflection based shims used by beam. Reflection is convenient for development,
// but is an unnecessary expense on pipeline performance.
//
// Using This Tool
//
// This tool is intended for use with `go generate`. The recommended convention
// putting the types and functions used in a separate package from pipeline construction.
// Then, the tool can be used as follows:
//
// //go:generate go install github.com/apache/beam/sdks/go/cmd/starcgen
// //go:generate starcgen --package=<mypackagename>
// //go:generate go fmt
//
// This will generate registrations and shim types for all types and functions
// in the package, in a file `<mypackagename>.shims.go`.
//
// Alternatively, it's possible to specify the specific input files and identifiers within
// the package for generation.
//
// //go:generate go install github.com/apache/beam/sdks/go/cmd/starcgen
// //go:generate starcgen --package=<mypackagename> --inputs=foo.go --identifiers=myFn,myStructFn --output=custom.shims.go
// //go:generate go fmt
//
package main
import (
"flag"
"fmt"
"go/ast"
"go/importer"
"go/parser"
"go/token"
"io"
"log"
"os"
"path/filepath"
"strings"
"github.com/apache/beam/sdks/go/pkg/beam/util/starcgenx"
)
var (
inputs = flag.String("inputs", "", "comma separated list of file with types for which to create shims")
intendedPkg = flag.String("package", "", "a filter on input go files. Required if inputs unset.")
output = flag.String("output", "", "output file with types to create")
ids = flag.String("identifiers", "", "comma separated list of package local identifiers for which to generate code")
debug = flag.Bool("debug", false, "print out a debugging header in the shim file to help diagnose errors")
)
// Generate takes the typechecked inputs, and generates the shim file for the relevant
// identifiers.
func Generate(w io.Writer, filename, pkg string, ids []string, fset *token.FileSet, files []*ast.File) error {
e := starcgenx.NewExtractor(pkg)
e.Ids = ids
e.Debug = *debug
// Importing from source should work in most cases.
imp := importer.For("source", nil)
if err := e.FromAsts(imp, fset, files); err != nil {
// Always print out the debugging info to the file.
if _, errw := w.Write(e.Bytes()); errw != nil {
return fmt.Errorf("error writing debug data to file after err %v:%v", err, errw)
}
err = fmt.Errorf("error extracting from asts: %v", err)
e.Printf("%v", err)
return err
}
data := e.Generate(filename)
if err := write(w, []byte(license)); err != nil {
return err
}
return write(w, data)
}
func write(w io.Writer, data []byte) error {
n, err := w.Write(data)
if err != nil && n < len(data) {
return fmt.Errorf("short write of data got %d, want %d", n, len(data))
}
return err
}
func usage() {
fmt.Fprintf(os.Stderr, "Usage: %v [options] --inputs=<comma separated of go files>\n", filepath.Base(os.Args[0]))
flag.PrintDefaults()
}
func main() {
flag.Usage = usage
flag.Parse()
log.SetFlags(log.Lshortfile)
log.SetPrefix("starcgen: ")
dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
if err != nil {
log.Fatal(err)
}
var ipts []string
// If inputs are empty, parse all go files in the local directory.
if len(*inputs) == 0 {
if *intendedPkg == "" {
log.Fatal("--package flag is required to be set when --inputs unset")
}
globbed, err := filepath.Glob(filepath.Join(dir, "*.go"))
if err != nil {
log.Fatal(err)
}
ipts = globbed
} else {
ipts = strings.Split(*inputs, ",")
}
// Get an output file for pre-processing if necessary.
if *output == "" && *intendedPkg != "" {
*output = *intendedPkg + ".shims.go"
}
outputBase := filepath.Base(*output)
fset := token.NewFileSet()
var fs []*ast.File
pkg := *intendedPkg
for _, i := range ipts {
// Ignore the existing shim file when re-generating.
if strings.HasSuffix(i, outputBase) {
continue
}
f, err := parser.ParseFile(fset, i, nil, 0)
if err != nil {
log.Fatal(err) // parse error
}
if pkg == "" {
pkg = f.Name.Name
} else if pkg != f.Name.Name {
// If we've set an intended package, just filter files outside that package.
// eg. for pkg_tests in the same directory.
if *intendedPkg != "" {
continue
}
log.Fatalf("Input file %v has mismatched package path, got %q, want %q", i, f.Name.Name, pkg)
}
fs = append(fs, f)
}
if pkg == "" {
log.Fatalf("No package detected in input files: %v", inputs)
}
f, err := os.OpenFile(*output, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
if err != nil {
log.Fatalf("error opening %q: %v", *output, err)
}
splitIds := make([]string, 0) // If no ids are specified, we should pass an empty slice.
if len(*ids) > 0 {
splitIds = strings.Split(*ids, ",")
}
if err := Generate(f, *output, pkg, splitIds, fset, fs); err != nil {
log.Fatal(err)
}
}
const license = `// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
`