forked from infobloxopen/protoc-gen-gorm
/
imports.go
120 lines (106 loc) · 3.67 KB
/
imports.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 plugin
import (
"fmt"
"sort"
"strings"
"github.com/gogo/protobuf/protoc-gen-gogo/generator"
)
/* --------- Response file import cleaning -------- */
// Imports that are added by default but unneeded in GORM code
var unneededImports = []string{
"import proto \"github.com/gogo/protobuf/proto\"\n",
"import _ \"github.com/infobloxopen/protoc-gen-gorm/options\"\n",
// if needed will be imported with an alias
"import _ \"github.com/infobloxopen/protoc-gen-gorm/types\"\n",
"var _ = proto.Marshal\n",
}
// CleanImports removes extraneous imports and lines from a proto response
// file Content
func CleanImports(pFileText *string) *string {
if pFileText == nil {
return pFileText
}
fileText := *pFileText
for _, dep := range unneededImports {
fileText = strings.Replace(fileText, dep, "", -1)
}
return &fileText
}
/* --------- Plugin level import handling --------- */
var (
gormImport = "github.com/jinzhu/gorm"
tkgormImport = "github.com/infobloxopen/atlas-app-toolkit/gorm"
uuidImport = "github.com/satori/go.uuid"
authImport = "github.com/infobloxopen/atlas-app-toolkit/auth"
gormpqImport = "github.com/jinzhu/gorm/dialects/postgres"
gtypesImport = "github.com/infobloxopen/protoc-gen-gorm/types"
ptypesImport = "github.com/golang/protobuf/ptypes"
wktImport = "github.com/golang/protobuf/ptypes/wrappers"
resourceImport = "github.com/infobloxopen/atlas-app-toolkit/gorm/resource"
fmImport = "google.golang.org/genproto/protobuf/field_mask"
queryImport = "github.com/infobloxopen/atlas-app-toolkit/query"
gatewayImport = "github.com/infobloxopen/atlas-app-toolkit/gateway"
)
type pkgImport struct {
packagePath string
alias string
}
// Import takes a package and adds it to the list of packages to import
// It will generate a unique new alias using the last portion of the import path
// unless the package is already imported for this file. Either way, it returns
// the package alias
func (p *OrmPlugin) Import(packagePath string) string {
subpath := packagePath[strings.LastIndex(packagePath, "/")+1:]
// package will always be suffixed with an integer to prevent any collisions
// with standard package imports
for i := 1; ; i++ {
newAlias := fmt.Sprintf("%s%d", strings.Replace(subpath, ".", "_", -1), i)
if pkg, ok := p.GetFileImports().packages[newAlias]; ok {
if packagePath == pkg.packagePath {
return pkg.alias
}
} else {
p.GetFileImports().packages[newAlias] = &pkgImport{packagePath: packagePath, alias: newAlias}
return newAlias
}
}
// Should never reach here
}
// UsingGoImports should be used with basic packages like "time", or "context"
func (p *OrmPlugin) UsingGoImports(pkgNames ...string) {
p.GetFileImports().stdImports = append(p.GetFileImports().stdImports, pkgNames...)
}
type fileImports struct {
wktPkgName string
typesToRegister []string
stdImports []string
packages map[string]*pkgImport
}
func newFileImports() *fileImports {
return &fileImports{packages: make(map[string]*pkgImport)}
}
func (p *OrmPlugin) GetFileImports() *fileImports {
return p.fileImports[p.currentFile]
}
// GenerateImports writes out required imports for the generated files
func (p *OrmPlugin) GenerateImports(file *generator.FileDescriptor) {
imports := p.fileImports[file]
for _, typeName := range imports.typesToRegister {
p.RecordTypeUse(typeName)
}
githubImports := imports.packages
sort.Strings(imports.stdImports)
for _, dep := range imports.stdImports {
p.PrintImport(dep, dep)
}
p.P()
aliases := []string{}
for a := range githubImports {
aliases = append(aliases, a)
}
sort.Strings(aliases)
for _, a := range aliases {
p.PrintImport(a, githubImports[a].packagePath)
}
p.P()
}