-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
120 lines (106 loc) · 3 KB
/
main.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
package main
import (
"flag"
"fmt"
"go/ast"
"go/parser"
"go/token"
"io/ioutil"
"log"
"os"
"strings"
"unicode"
)
var (
inputFile = flag.String("in", "", "input file to generate type registry from")
outputFile = flag.String("out", "", "output file name; default srcdir/<pbfile>.typeregistry.go")
packageName = flag.String("package", "", "Package name to use; default retrieved from input file")
fieldName = flag.String("registryfield", "typeregistry", "Field name to declare the registry under")
typenamePrefix = flag.String("typenameprefix", "", "Prefix to add to the type name values emitted")
)
func Usage() {
fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
fmt.Fprint(os.Stderr, "\ttype-registry -in mytypes.pb.go\n")
fmt.Fprint(os.Stderr, "Flags:\n")
flag.PrintDefaults()
}
func main() {
log.SetFlags(0)
log.SetPrefix("type-registry: ")
flag.Usage = Usage
flag.Parse()
if len(*inputFile) == 0 {
flag.Usage()
os.Exit(2)
}
// read in .pb.go file
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, *inputFile, nil, 0)
if err != nil {
panic(err)
}
g := NewGenerator()
// determine package name
if len(*packageName) == 0 {
g.packageName = f.Name.Name
} else {
g.packageName = *packageName
}
g.innerFieldName = *fieldName
g.typenamePrefix = *typenamePrefix
// find all the type specs and note them for later
ast.Inspect(f, func(n ast.Node) bool {
switch x := n.(type) {
case *ast.TypeSpec:
switch x.Type.(type) {
case *ast.StructType:
s := x.Name.Name
// exported type name?
if unicode.IsUpper([]rune(s)[0]) {
g.typeNames[s] = true
}
}
}
return true
})
// write to the buf in generator
g.Printf("// Code generated by \"type-registry %s\"; DO NOT EDIT.\n", strings.Join(os.Args[1:], " "))
g.Printf("\n")
g.Printf("package %s", g.packageName)
g.Printf("\n")
// imports
g.Printf("import \"reflect\"\n")
g.Printf("\n")
g.Printf("var %s = make(map[string]reflect.Type)\n", g.innerFieldName)
g.Printf("\n")
g.Printf("func init() {\n")
// generated code needs to be stable to avoid churning source control
for _, t := range g.SortedTypeNames() {
g.Printf("%s[\"%s\"] = reflect.TypeOf(%s{})\n", g.innerFieldName, t, g.transformTypeName(t))
}
g.Printf("}")
src := g.format()
// write to file
outputName := determineFileName(*inputFile, *outputFile)
if err := ioutil.WriteFile(outputName, src, 0644); err != nil {
log.Fatalf("writing output: %s", err)
}
}
// Figure out what file to write to.
// If override is given, just use that.
// Otherwise determine default from inputFileName,
// where "path/foo.pb.go" becomes "path/foo.pb.typeregistry.go"
func determineFileName(inputFileName string, override string) string {
if override != "" {
return override
}
parts := strings.Split(inputFileName, ".")
// is the last part ".go"?
if parts[len(parts)-1] == "go" {
// remove the extra .go
parts = parts[0 : len(parts)-1]
}
// add the new suffix
parts = append(parts, "typeregistry", "go")
return strings.Join(parts, ".")
}