package gobuf import ( "fmt" "go/ast" "go/parser" "go/token" "io" "log" "os" "path/filepath" "regexp" "strings" "unicode" ) var ( rComment = regexp.MustCompile(`^//.*?@(?i:gotags?):\s*(.*)$`) bsonComment = regexp.MustCompile(`^//.*?@(?i:bson?):\s*(.*)$`) jsonComment = regexp.MustCompile(`^//.*?@(?i:json?):\s*(.*)$`) rInject = regexp.MustCompile("`.+`$") rTags = regexp.MustCompile(`[\w_]+:"[^"]+"`) ) type textArea struct { Start int End int CurrentTag string InjectTag string CommentStart int CommentEnd int } type TagValueStyle int const ( Underline TagValueStyle = iota LowerCase UpperCase ) type InjectTagProps struct { TagName string // tag name Style TagValueStyle // tag value style, default underline comment *ast.Comment fieldName string fieldValue string } type InjectTag struct { inputFiles string defaultTags map[string]TagValueStyle } func NewInjectTag(pbFiles string) *InjectTag { return &InjectTag{inputFiles: pbFiles} } func (it *InjectTag) Inject() error { globResults, err := filepath.Glob(it.inputFiles) if err != nil { return fmt.Errorf("parser input file failed: %v", err) } var matched int for _, path := range globResults { fileInfo, err := os.Stat(path) if err != nil { return fmt.Errorf("read file stat failed: %v", err) } if fileInfo.IsDir() { continue } // It should end with ".go" at a minimum. if !strings.HasSuffix(strings.ToLower(fileInfo.Name()), ".go") { continue } matched++ areas, err := it.parserFile(path) if err != nil { log.Fatal(err) } if err = it.writeFile(path, areas); err != nil { log.Fatal(err) } } if matched == 0 { return fmt.Errorf("input %q matched no files, see: -help", it.inputFiles) } return nil } func (it *InjectTag) WithTags(tags ...InjectTagProps) *InjectTag { if len(it.defaultTags) == 0 { it.defaultTags = make(map[string]TagValueStyle) } for _, tag := range tags { it.defaultTags[tag.TagName] = tag.Style } return it } func (it *InjectTag) parserFile(inputPath string) (areas []textArea, err error) { it.logf("parsing file %q for inject tag comments", inputPath) fset := token.NewFileSet() f, err := parser.ParseFile(fset, inputPath, nil, parser.ParseComments) if err != nil { return } for _, decl := range f.Decls { // check if is generic declaration genDecl, ok := decl.(*ast.GenDecl) if !ok { continue } var typeSpec *ast.TypeSpec for _, spec := range genDecl.Specs { if ts, tsOK := spec.(*ast.TypeSpec); tsOK { typeSpec = ts break } } // skip if can't get type spec if typeSpec == nil { continue } // not a struct, skip structDecl, ok := typeSpec.Type.(*ast.StructType) if !ok { continue } for _, field := range structDecl.Fields.List { // skip if field name abnormal if len(field.Names) != 1 { continue } fieldName := field.Names[0].Name if unicode.IsLower(rune(fieldName[0])) { continue } // skip if field has no doc var comments []*ast.Comment if field.Doc != nil { comments = append(comments, field.Doc.List...) } // The "doc" field (above comment) is more commonly "free-form" // due to the ability to have a much larger comment without it // being unwieldy. As such, the "comment" field (trailing comment), // should take precedence if there happen to be multiple tags // specified, both in the field doc, and the field line. Whichever // comes last, will take precedence. if field.Comment != nil { comments = append(comments, field.Comment.List...) } tags := it.tagFromComment(field.Names[0].Name, comments) if len(tags) == 0 { continue } currentTag := field.Tag.Value for _, tag := range tags { area := textArea{ Start: int(field.Pos()), End: int(field.End()), CurrentTag: currentTag[1 : len(currentTag)-1], InjectTag: tag.fieldValue, } if tag.comment != nil { area.CommentStart = int(tag.comment.Pos()) area.CommentEnd = int(tag.comment.End()) } areas = append(areas, area) } } } //it.logf("parsed file %q, number of fields to inject custom tags: %d", inputPath, len(areas)) return } func (it *InjectTag) writeFile(inputPath string, areas []textArea) error { f, err := os.Open(inputPath) if err != nil { return err } contents, err := io.ReadAll(f) if err != nil { return err } if err = f.Close(); err != nil { return err } // inject custom tags from tail of file first to preserve order for i := range areas { area := areas[len(areas)-i-1] it.logf("inject custom tag %q to expression %q", area.InjectTag, string(contents[area.Start-1:area.End-1])) contents = it.injectTag(contents, area) } if err = os.WriteFile(inputPath, contents, 0o644); err != nil { return err } if len(areas) > 0 { it.logf("file %q is injected with custom tags", inputPath) } return nil } func (it *InjectTag) logf(format string, v ...interface{}) { log.Printf(format, v...) } func (it *InjectTag) tagFromComment(fieldName string, comments []*ast.Comment) (tags []InjectTagProps) { var commentInject = make(map[string]struct{}) for i := range comments { bsonMatch := bsonComment.FindStringSubmatch(comments[i].Text) if len(bsonMatch) == 2 { tags = append(tags, InjectTagProps{ TagName: "bson", comment: comments[i], fieldName: fieldName, fieldValue: fmt.Sprintf(`bson:%v`, bsonMatch[1]), }) commentInject["bson"] = struct{}{} continue } jsonMatch := jsonComment.FindStringSubmatch(comments[i].Text) if len(jsonMatch) == 2 { tags = append(tags, InjectTagProps{ TagName: "json", comment: comments[i], fieldName: fieldName, fieldValue: fmt.Sprintf(`json:%v`, jsonMatch[1]), }) commentInject["json"] = struct{}{} continue } match := rComment.FindStringSubmatch(comments[i].Text) if len(match) == 2 { tagVal := match[1] tagName := strings.Split(match[1], ":")[0] tags = append(tags, InjectTagProps{ TagName: tagName, comment: comments[i], fieldName: fieldName, fieldValue: tagVal, }) commentInject[tagName] = struct{}{} continue } } for tag, style := range it.defaultTags { _, exist := commentInject[tag] if exist { continue } tags = append(tags, InjectTagProps{ TagName: tag, fieldName: fieldName, fieldValue: fmt.Sprintf("%s:\"%s\"", tag, getFieldValue(fieldName, style)), }) } return } type tagItem struct { key string value string } type tagItems []tagItem func (ti *tagItems) format() string { var tags []string for _, item := range *ti { tags = append(tags, fmt.Sprintf(`%s:%s`, item.key, item.value)) } return strings.Join(tags, " ") } func (ti *tagItems) override(nti tagItems) tagItems { var overrides []tagItem for i := range *ti { dup := -1 for j := range nti { if (*ti)[i].key == nti[j].key { dup = j break } } if dup == -1 { overrides = append(overrides, (*ti)[i]) } else { overrides = append(overrides, nti[dup]) nti = append(nti[:dup], nti[dup+1:]...) } } return append(overrides, nti...) } func (it *InjectTag) newTagItems(tag string) tagItems { var items []tagItem it.logf("new tag: %v\n", tag) split := rTags.FindAllString(tag, -1) for _, t := range split { sepPos := strings.Index(t, ":") items = append(items, tagItem{ key: t[:sepPos], value: t[sepPos+1:], }) } return items } func (it *InjectTag) injectTag(contents []byte, area textArea) (injected []byte) { expr := make([]byte, area.End-area.Start) copy(expr, contents[area.Start-1:area.End-1]) log.Println("expr: ", string(expr)) cti := it.newTagItems(area.CurrentTag) log.Printf("cti: %v\n", cti) iti := it.newTagItems(area.InjectTag) log.Printf("iti: %v\n", iti) ti := cti.override(iti) log.Printf("ti: %v\n", ti) expr = rInject.ReplaceAll(expr, []byte(fmt.Sprintf("`%s`", ti.format()))) injected = append(injected, contents[:area.Start-1]...) injected = append(injected, expr...) injected = append(injected, contents[area.End-1:]...) return }