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  }