go generate and ast

//go:generate stringer -type=Pill

package painkiller

type Pill int

const (
    Placebo Pill = iota
    Aspirin
    Ibuprofen
    Paracetamol
    Acetaminophen = Paracetamol
)

// generated by stringer -type Pill pill.go; DO NOT EDIT

package painkiller

import "fmt"

const _Pill_name = "PlaceboAspirinIbuprofenParacetamol"

var _Pill_index = [...]uint8{0, 7, 14, 23, 34}

func (i Pill) String() string {
    if i < 0 || i+1 >= Pill(len(_Pill_index)) {
        return fmt.Sprintf("Pill(%d)", i)
    }
    return _Pill_name[_Pill_index[i]:_Pill_index[i+1]]
}

// Node AST树节点
type Node interface {
  Pos() token.Pos
  End() token.Pos
}

// Expr 所有的表达式都需要实现Expr接口
type Expr interface {
  Node
  exprNode()
}

// Stmt 所有的语句都需要实现Stmt接口
type Stmt interface {
  Node
  stmtNode()
}

// Decl 所有的声明都需要实现Decl接口
type Decl interface {
  Node
  declNode()
}

// ValueSpec 表示常量声明或者变量声明
type ValueSpec struct {
        Doc     *CommentGroup // associated documentation; or nil
        Names   []*Ident      // value names (len(Names) > 0)
        Type    Expr          // value type; or nil
        Values  []Expr        // initial values; or nil
        Comment *CommentGroup // line comments; or nil
}

// CommentMap把AST节点和其关联的注释列表进行映射
type CommentMap map[Node][]*CommentGroup

file := os.Getenv("GOFILE")
// 保存注释信息
var comments = make(map[string]string)

// 解析代码源文件，获取常量和注释之间的关系
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
checkErr(err)

// Create an ast.CommentMap from the ast.File's comments.
// This helps keeping the association between comments
// and AST nodes.
cmap := ast.NewCommentMap(fset, f, f.Comments)
for node := range cmap {
  // 仅支持一条声明语句，一个常量的情况
  if spec, ok := node.(*ast.ValueSpec); ok && len(spec.Names) == 1 {
    // 仅提取常量的注释
    ident := spec.Names[0]
    if ident.Obj.Kind == ast.Con {
      // 获取注释信息
      comments[ident.Name] = getComment(ident.Name, spec.Doc)
    }
  }
}

// getComment 获取注释信息，来自AST标准库的summary方法
func getComment(name string, group *ast.CommentGroup) string {
	var buf bytes.Buffer


	for _, comment := range group.List {
    // 注释信息会以 // 参数名，开始，我们实际使用时不需要，去掉
		text := strings.TrimSpace(strings.TrimPrefix(comment.Text, fmt.Sprintf("// %s", name)))
		buf.WriteString(text)
	}

	// replace any invisibles with blanks
	bytes := buf.Bytes()
	for i, b := range bytes {
		switch b {
		case '\t', '\n', '\r':
			bytes[i] = ' '
		}
	}

	return string(bytes)
}

const suffix = "_msg_gen.go"

// tpl 生成代码需要用到模板
const tpl = `
// Code generated by github.com/mohuishou/gen-const-msg DO NOT EDIT

// {{.pkg}} const code comment msg
package {{.pkg}}

// noErrorMsg if code is not found, GetMsg will return this
const noErrorMsg = "unknown error"

// messages get msg from const comment
var messages = map[int]string{
	{{range $key, $value := .comments}}
	{{$key}}: "{{$value}}",{{end}}
}

// GetMsg get error msg
func GetMsg(code int) string {
	var (
		msg string
		ok  bool
	)
	if msg, ok = messages[code]; !ok {
		msg = noErrorMsg
	}
	return msg
}
`

// gen 生成代码
func gen(comments map[string]string) ([]byte, error) {
  var buf = bytes.NewBufferString("")

  data := map[string]interface{}{
    "pkg":      os.Getenv("GOPACKAGE"),
    "comments": comments,
  }

  t, err := template.New("").Parse(tpl)
  if err != nil {
    return nil, errors.Wrapf(err, "template init err")
  }

  err = t.Execute(buf, data)
  if err != nil {
    return nil, errors.Wrapf(err, "template data err")
  }

  return format.Source(buf.Bytes())
}