@@ -11,10 +11,20 @@ import (
1111 "log"
1212 "path/filepath"
1313 "strings"
14+ "time"
1415
1516 "golang.org/x/tools/imports"
1617)
1718
19+ type makeInterfaceFile struct {
20+ DirPath string
21+ PkgName string
22+ Structs []string
23+ TypeDoc map [string ]string
24+ AllMethods map [string ][]string
25+ AllImports []string
26+ }
27+
1828type Method struct {
1929 Code string
2030 Docs []string
@@ -50,7 +60,7 @@ func getReceiverType(fd *ast.FuncDecl) (ast.Expr, error) {
5060 return fd .Recv .List [0 ].Type , nil
5161}
5262
53- func formatFieldList (src []byte , fl * ast.FieldList , pkgName string ) []string {
63+ func formatFieldList (src []byte , fl * ast.FieldList ) []string {
5464 if fl == nil {
5565 return nil
5666 }
@@ -72,44 +82,6 @@ func formatFieldList(src []byte, fl *ast.FieldList, pkgName string) []string {
7282 return parts
7383}
7484
75- func formatCode (code string ) ([]byte , error ) {
76- opts := & imports.Options {
77- TabIndent : true ,
78- TabWidth : 2 ,
79- Fragment : true ,
80- Comments : true ,
81- }
82- return imports .Process ("" , []byte (code ), opts )
83- }
84-
85- func makeInterface (pkgName string , ifaceComment map [string ]string , structName string , methods []string , imports []string ) ([]byte , error ) {
86- output := []string {
87- "// Code generated by struct2interface; DO NOT EDIT." ,
88- "" ,
89- "package " + pkgName ,
90- "import (" ,
91- }
92- output = append (output , imports ... )
93- output = append (output ,
94- ")" ,
95- "" ,
96- )
97-
98- output = append (output , fmt .Sprintf ("// %s" , strings .Replace (ifaceComment [structName ], "\n " , "\n // " , - 1 )))
99- output = append (output , fmt .Sprintf ("type %s interface {" , structName + "Interface" ))
100- output = append (output , methods ... )
101- output = append (output , "}" )
102-
103- code := strings .Join (output , "\n " )
104-
105- result , err := formatCode (code )
106- if err != nil {
107- fmt .Println (code )
108- return nil , err
109- }
110- return result , nil
111- }
112-
11385func parseStruct (src []byte ) (pkgName string , structs []string , methods map [string ][]Method , imports []string , typeDoc map [string ]string , err error ) {
11486 fset := token .NewFileSet ()
11587 a , err := parser .ParseFile (fset , "" , src , parser .ParseComments )
@@ -129,68 +101,156 @@ func parseStruct(src []byte) (pkgName string, structs []string, methods map[stri
129101
130102 methods = make (map [string ][]Method )
131103 for _ , d := range a .Decls {
132- if a , fd := getReceiverTypeName (src , d ); a != "" {
104+ if structName , fd := getReceiverTypeName (src , d ); structName != "" {
105+ // 私有方法
133106 if ! fd .Name .IsExported () {
134107 continue
135108 }
136- params := formatFieldList (src , fd .Type .Params , pkgName )
137- ret := formatFieldList (src , fd .Type .Results , pkgName )
109+ params := formatFieldList (src , fd .Type .Params )
110+ ret := formatFieldList (src , fd .Type .Results )
138111 method := fmt .Sprintf ("%s(%s) (%s)" , fd .Name .String (), strings .Join (params , ", " ), strings .Join (ret , ", " ))
139112 var docs []string
140113 if fd .Doc != nil {
141114 for _ , d := range fd .Doc .List {
142115 docs = append (docs , string (src [d .Pos ()- 1 :d .End ()- 1 ]))
143116 }
144117 }
145- if _ , ok := methods [a ]; ! ok {
146- structs = append (structs , a )
118+ if _ , ok := methods [structName ]; ! ok {
119+ structs = append (structs , structName )
147120 }
148121
149- methods [a ] = append (methods [a ], Method {
122+ methods [structName ] = append (methods [structName ], Method {
150123 Code : method ,
151124 Docs : docs ,
152125 })
153126 }
154127 }
155128
156129 typeDoc = make (map [string ]string )
157- pkg := & ast.Package {Files : map [string ]* ast.File {"" : a }}
158- doc := doc .New (pkg , "" , doc .AllDecls )
159- for _ , t := range doc .Types {
130+ for _ , t := range doc .New (& ast.Package {Files : map [string ]* ast.File {"" : a }}, "" , doc .AllDecls ).Types {
160131 typeDoc [t .Name ] = strings .TrimSuffix (t .Doc , "\n " )
161132 }
162133
163134 return
164135}
165136
166- func makeFile (file string ) ([]byte , error ) {
167- allMethods := make (map [string ][]string )
168- allImports := []string {}
169- var structs []string
170- // mset := make(map[string]struct{})
171- iset := make (map [string ]struct {})
172- typeDoc := make (map [string ]string )
173- pkgName := ""
137+ func formatCode (code string ) ([]byte , error ) {
138+ opts := & imports.Options {
139+ TabIndent : true ,
140+ TabWidth : 2 ,
141+ Fragment : true ,
142+ Comments : true ,
143+ }
144+ return imports .Process ("" , []byte (code ), opts )
145+ }
146+
147+ func makeInterfaceHead (pkgName string , imports []string ) []string {
148+ output := []string {
149+ "// Code generated by struct2interface; DO NOT EDIT." ,
150+ "" ,
151+ "package " + pkgName ,
152+ "import (" ,
153+ }
154+ output = append (output , imports ... )
155+ output = append (output ,
156+ ")" ,
157+ "" ,
158+ )
159+ return output
160+ }
161+
162+ func makeInterfaceBody (output []string , ifaceComment map [string ]string , structName string , methods []string ) []string {
163+
164+ comment := strings .TrimSuffix (strings .Replace (ifaceComment [structName ], "\n " , "\n //\t " , - 1 ), "\n //\t " )
165+ if len (strings .TrimSpace (comment )) > 0 {
166+ output = append (output , fmt .Sprintf ("// %s" , comment ))
167+ }
168+
169+ output = append (output , fmt .Sprintf ("type %s interface {" , structName + "Interface" ))
170+ output = append (output , methods ... )
171+ output = append (output , "}" )
172+ return output
173+ }
174+
175+ func createFile (objs map [string ][]* makeInterfaceFile ) error {
176+ for dir , obj := range objs {
177+ if len (obj ) == 0 {
178+ continue
179+ }
180+
181+ var (
182+ startTime = time .Now ()
183+ firstObj = obj [0 ]
184+ pkgName = firstObj .PkgName
185+ typeDoc = firstObj .TypeDoc
186+ mapStructMethods = make (map [string ][]string )
187+ listStructMethods = make ([]string , 0 )
188+ structAllImports = make ([]string , 0 )
189+ )
190+
191+ for _ , file := range obj {
192+ for _ , structName := range file .Structs {
193+ if _ , ok := mapStructMethods [structName ]; ok {
194+ mapStructMethods [structName ] = append (mapStructMethods [structName ], file .AllMethods [structName ]... )
195+ } else {
196+ mapStructMethods [structName ] = file .AllMethods [structName ]
197+ listStructMethods = append (listStructMethods , structName )
198+ }
199+
200+ structAllImports = append (structAllImports , file .AllImports ... )
201+ }
202+ }
203+
204+ output := makeInterfaceHead (pkgName , structAllImports )
205+
206+ for _ , structName := range listStructMethods {
207+ methods , ok := mapStructMethods [structName ]
208+ if ! ok {
209+ continue
210+ }
211+ output = makeInterfaceBody (output , typeDoc , structName , methods )
212+ }
213+
214+ code := strings .Join (output , "\n " )
215+ result , err := formatCode (code )
216+ if err != nil {
217+ fmt .Printf ("[struct2interface] %s \n " , "formatCode error" )
218+ return err
219+ }
220+ var fileName = filepath .Join (dir , "interface_" + pkgName + ".go" )
221+ if err = ioutil .WriteFile (fileName , result , 0644 ); err != nil {
222+ return err
223+ }
224+ fmt .Printf ("[struct2interface] %s %s %s \n " , "parsing" , time .Since (startTime ).String (), fileName )
225+ }
226+
227+ return nil
228+ }
229+
230+ func makeFile (file string ) (* makeInterfaceFile , error ) {
231+ var (
232+ allMethods = make (map [string ][]string )
233+ allImports = make ([]string , 0 )
234+ iset = make (map [string ]struct {})
235+ typeDoc = make (map [string ]string )
236+ )
174237
175238 src , err := ioutil .ReadFile (file )
176239 if err != nil {
177240 return nil , err
178241 }
179242
180- pkg , structSlice , methods , imports , parsedTypeDoc , err := parseStruct (src )
243+ pkgName , structSlice , methods , importList , parsedTypeDoc , err := parseStruct (src )
181244 if err != nil {
182- log . Println ( "file: "file )
245+ fmt . Printf ( "[struct2interface] %s, err: %s \n "" file parseStruct error" , err . Error () )
183246 return nil , err
184247 }
185248
186249 if len (methods ) == 0 {
187250 return nil , nil
188251 }
189252
190- pkgName = pkg
191- structs = structSlice
192-
193- for _ , i := range imports {
253+ for _ , i := range importList {
194254 if _ , ok := iset [i ]; ! ok {
195255 allImports = append (allImports , i )
196256 iset [i ] = struct {}{}
@@ -199,48 +259,30 @@ func makeFile(file string) ([]byte, error) {
199259
200260 for structName , mm := range methods {
201261 typeDoc [structName ] = fmt .Sprintf ("%s ...\n %s" , structName + "Interface" , parsedTypeDoc [structName ])
202-
203262 for _ , m := range mm {
204-
205- // if _, ok := mset[m.Code]; !ok {
206263 allMethods [structName ] = append (allMethods [structName ], m .Lines ()... )
207- // mset[m.Code] = struct{}{}
208- // }
209- }
210- }
211-
212- var result []byte
213-
214- for _ , structName := range structs {
215- result , err = makeInterface (pkgName , typeDoc , structName , allMethods [structName ], allImports )
216- if err != nil {
217- return nil , err
218264 }
219-
220- dir := filepath .Dir (file )
221- output := filepath .Join (dir , "interface_" + structName + ".go" )
222-
223- err := ioutil .WriteFile (output , result , 0644 )
224- if err != nil {
225- return nil , err
226- }
227- fmt .Println ("struct2interface:" , dir + ": wrote" , output )
228265 }
229266
230- return result , nil
267+ return & makeInterfaceFile {
268+ DirPath : filepath .Dir (file ),
269+ PkgName : pkgName ,
270+ Structs : structSlice ,
271+ TypeDoc : typeDoc ,
272+ AllMethods : allMethods ,
273+ AllImports : allImports ,
274+ }, nil
231275}
232276
233277func MakeDir (dir string ) error {
234-
235- err := filepath .WalkDir (dir , func (path string , d fs.DirEntry , err error ) error {
278+ var mapDirPath = make ( map [ string ][] * makeInterfaceFile )
279+ if err := filepath .WalkDir (dir , func (path string , d fs.DirEntry , err error ) error {
236280 if err != nil {
237- log . Fatal ( err )
281+ return err
238282 }
239-
240283 if d .IsDir () {
241284 return nil
242285 }
243-
244286 if strings .HasPrefix (filepath .Base (path ), "interface_" ) {
245287 return nil
246288 }
@@ -253,18 +295,22 @@ func MakeDir(dir string) error {
253295
254296 result , err := makeFile (path )
255297 if err != nil {
256- log .Fatal ("struct2interface.Make failed," , err .Error (), path )
298+ log .Panic ("struct2interface.Make failed," , err .Error (), path )
299+ } else if result == nil {
300+ return nil
257301 }
258302
259- if len (result ) == 0 {
260- return nil
303+ if _ , ok := mapDirPath [filepath .Dir (path )]; ok {
304+ mapDirPath [filepath .Dir (path )] = append (mapDirPath [filepath .Dir (path )], result )
305+ } else {
306+ mapDirPath [filepath .Dir (path )] = []* makeInterfaceFile {result }
261307 }
262308
263309 return nil
264- })
265- if err != nil {
310+ }); err != nil {
311+ fmt . Printf ( "[struct2interface] %s \n " , err . Error ())
266312 return err
267313 }
268314
269- return nil
315+ return createFile ( mapDirPath )
270316}
0 commit comments