github.com/galaxyobe/gen@v0.0.0-20220910125335-392fa8f0990f/pkg/util/ast.go (about) 1 /* 2 Copyright 2022 Galaxyobe. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package util 18 19 import ( 20 "bytes" 21 "go/ast" 22 "go/parser" 23 "go/token" 24 "io/fs" 25 "path/filepath" 26 "strings" 27 ) 28 29 func walkAstParseFile(files *[]*ast.File) fs.WalkDirFunc { 30 return func(path string, d fs.DirEntry, err error) error { 31 if err != nil { 32 return err 33 } 34 if strings.HasSuffix(path, "_generated.go") { 35 return nil 36 } 37 if d.IsDir() { 38 walkAstParseFile(files) 39 return nil 40 } 41 fset := token.NewFileSet() 42 f, err := parser.ParseFile(fset, path, nil, parser.DeclarationErrors|parser.ParseComments) 43 if err != nil { 44 return err 45 } 46 *files = append(*files, f) 47 return nil 48 } 49 } 50 51 func AstParseDir(dir string) (files []*ast.File, err error) { 52 err = filepath.WalkDir(dir, walkAstParseFile(&files)) 53 return 54 } 55 56 func getAstObjectName(o any) string { 57 switch v := o.(type) { 58 case *ast.Ident: 59 return v.Name 60 default: 61 return "" 62 } 63 } 64 65 func getFuncType(funcType *ast.FuncType) string { 66 var buf bytes.Buffer 67 buf.WriteString("func(") 68 for idx, param := range funcType.Params.List { 69 buf.WriteString(getAstObjectName(param.Type)) 70 if idx < len(funcType.Params.List)-1 { 71 buf.WriteByte(',') 72 } 73 } 74 buf.WriteByte(')') 75 if len(funcType.Results.List) > 1 { 76 buf.WriteByte('(') 77 } 78 for idx, param := range funcType.Results.List { 79 buf.WriteString(getAstObjectName(param.Type)) 80 if idx < len(funcType.Params.List)-1 { 81 buf.WriteByte(',') 82 } 83 } 84 if len(funcType.Results.List) > 1 { 85 buf.WriteByte(')') 86 } 87 return buf.String() 88 } 89 90 func getTypeInTypeSpec(typeSpec *ast.TypeSpec) (type_ string, name string) { 91 name = typeSpec.Name.Name 92 switch t := typeSpec.Type.(type) { 93 case *ast.SelectorExpr: 94 type_ = getAstObjectName(t.X) + "." + getAstObjectName(t.Sel) 95 case *ast.Ident: 96 type_ = t.Name 97 case *ast.FuncType: 98 // type_ = getFuncType(t) 99 type_ = "func" 100 case *ast.StructType: 101 type_ = "struct" 102 case *ast.StarExpr: 103 type_ = getAstObjectName(t.X) 104 default: 105 return 106 } 107 return 108 } 109 110 func getTypesInGenDecl(genDecl *ast.GenDecl) map[string]string { 111 var m = make(map[string]string) 112 for _, spec := range genDecl.Specs { 113 switch s := spec.(type) { 114 case *ast.TypeSpec: 115 type_, name := getTypeInTypeSpec(s) 116 m[name] = type_ 117 } 118 } 119 return m 120 } 121 122 func FindTypeNames(files []*ast.File) (m map[string]string) { 123 m = make(map[string]string) 124 for _, file := range files { 125 ast.Inspect(file, func(node ast.Node) bool { 126 switch n := node.(type) { 127 case *ast.GenDecl: 128 result := getTypesInGenDecl(n) 129 for k, v := range result { 130 m[k] = v 131 } 132 } 133 return true 134 }) 135 } 136 return 137 }