github.com/unionj-cloud/go-doudou/v2@v2.3.5/toolkit/astutils/rewritejsontag.go (about)

     1  package astutils
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"github.com/pkg/errors"
     7  	"github.com/unionj-cloud/go-doudou/v2/toolkit/caller"
     8  	"go/ast"
     9  	"go/format"
    10  	"go/parser"
    11  	"go/token"
    12  	"golang.org/x/tools/go/ast/astutil"
    13  	"regexp"
    14  	"strings"
    15  	"unicode"
    16  )
    17  
    18  func isExport(field string) bool {
    19  	return unicode.IsUpper([]rune(field)[0])
    20  }
    21  
    22  func extractJsonPropName(tag string) string {
    23  	re := regexp.MustCompile(`json:"(.*?)"`)
    24  	if re.MatchString(tag) {
    25  		subs := re.FindAllStringSubmatch(tag, -1)
    26  		return strings.TrimSpace(strings.Split(subs[0][1], ",")[0])
    27  	}
    28  	return ""
    29  }
    30  
    31  func extractFormPropName(tag string) string {
    32  	re := regexp.MustCompile(`form:"(.*?)"`)
    33  	if re.MatchString(tag) {
    34  		subs := re.FindAllStringSubmatch(tag, -1)
    35  		return strings.TrimSpace(strings.Split(subs[0][1], ",")[0])
    36  	}
    37  	return ""
    38  }
    39  
    40  type RewriteTagConfig struct {
    41  	File        string
    42  	Omitempty   bool
    43  	ConvertFunc func(old string) string
    44  	Form        bool
    45  }
    46  
    47  // RewriteTag overwrites json tag by convert function and return formatted source code
    48  func RewriteTag(config RewriteTagConfig) (string, error) {
    49  	file, convert, omitempty, form := config.File, config.ConvertFunc, config.Omitempty, config.Form
    50  	fset := token.NewFileSet()
    51  	root, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
    52  	if err != nil {
    53  		return "", errors.Wrap(err, caller.NewCaller().String())
    54  	}
    55  	re := regexp.MustCompile(`json:"(.*?)"`)
    56  	reForm := regexp.MustCompile(`form:"(.*?)"`)
    57  	astutil.Apply(root, func(cursor *astutil.Cursor) bool {
    58  		return true
    59  	}, func(cursor *astutil.Cursor) bool {
    60  		structSpec, ok := cursor.Node().(*ast.StructType)
    61  		if !ok {
    62  			return true
    63  		}
    64  		for _, field := range structSpec.Fields.List {
    65  			if field.Names == nil {
    66  				continue
    67  			}
    68  			fname := field.Names[0].Name
    69  			if !isExport(fname) {
    70  				continue
    71  			}
    72  			tagValue := convert(field.Names[0].Name)
    73  			jsonTagValue := tagValue
    74  			if omitempty {
    75  				jsonTagValue += ",omitempty"
    76  			}
    77  			jsonTag := fmt.Sprintf(`json:"%s"`, jsonTagValue)
    78  
    79  			formTagValue := tagValue
    80  			if omitempty {
    81  				formTagValue += ",omitempty"
    82  			}
    83  			formTag := fmt.Sprintf(`form:"%s"`, formTagValue)
    84  
    85  			if field.Tag != nil {
    86  				if re.MatchString(field.Tag.Value) {
    87  					if extractJsonPropName(field.Tag.Value) != "-" {
    88  						field.Tag.Value = re.ReplaceAllLiteralString(field.Tag.Value, jsonTag)
    89  					}
    90  				} else {
    91  					lastindex := strings.LastIndex(field.Tag.Value, "`")
    92  					if lastindex < 0 {
    93  						panic(errors.New("call LastIndex() error"))
    94  					}
    95  					field.Tag.Value = field.Tag.Value[:lastindex] + fmt.Sprintf(" %s`", jsonTag)
    96  				}
    97  				if form {
    98  					if reForm.MatchString(field.Tag.Value) {
    99  						if extractFormPropName(field.Tag.Value) != "-" {
   100  							field.Tag.Value = reForm.ReplaceAllLiteralString(field.Tag.Value, formTag)
   101  						}
   102  					} else {
   103  						lastindex := strings.LastIndex(field.Tag.Value, "`")
   104  						if lastindex < 0 {
   105  							panic(errors.New("call LastIndex() error"))
   106  						}
   107  						field.Tag.Value = field.Tag.Value[:lastindex] + fmt.Sprintf(" %s`", formTag)
   108  					}
   109  				}
   110  			} else {
   111  				if form {
   112  					field.Tag = &ast.BasicLit{
   113  						Kind:  token.STRING,
   114  						Value: fmt.Sprintf("`%s %s`", jsonTag, formTag),
   115  					}
   116  				} else {
   117  					field.Tag = &ast.BasicLit{
   118  						Kind:  token.STRING,
   119  						Value: fmt.Sprintf("`%s`", jsonTag),
   120  					}
   121  				}
   122  			}
   123  		}
   124  		return true
   125  	})
   126  	buf := &bytes.Buffer{}
   127  	err = format.Node(buf, fset, root)
   128  	if err != nil {
   129  		return "", fmt.Errorf("error formatting new code: %w", err)
   130  	}
   131  	return buf.String(), nil
   132  }