codemsg模块可以把感兴趣的模块提取到map里面

codemsg.md

@@ -3,6 +3,13 @@
 只要实现code，自动生成String()类型和CodeMsg类型。可以在错误码代码节约时间  
 1.给每个ErrNo类型生成String()方法, 内容就是注释里面的。  
 2.给每个ErrNo类型生成CodeMsg{}结构
+
+### 一、生成code msg代码
+1.1 使用命令
+```go
+h2o codemsg --code-msg --linecomment --type ErrNo ./testdata/err.go
+```
+1.2 code msg定义例子
 ```go
 // 这段代码是我们要写的。
 package demo
@@ -85,3 +92,33 @@ var (
 // 如果要生成grpc的错误就用下面的命令
 // h2o codemsg --code-msg --linecomment --type ErrNo ./testdata/err.go --grpc --string-method string2 --string
 ```
+
+### 二、提取某些code到指定的map里面 
+主要是过滤一些code，比如对错误进行分级的使用场景, 有些错误非服务端引起。需要忽略这些错误
+2.1 命令
+```go
+h2o codemsg --code-msg --linecomment --type ErrNo --take-code-to-map ./testdata/err.go
+```
+
+2.2 code msg定义例子, 提取感兴趣的code到map里面
+```go
+// 这段代码是我们要写的。
+package demo
+
+type ErrNo int32 // 这里的类型可以自定义，--type后面类型
+
+const (
+	ENo ErrNo = 1003 // @TakeCodeToMap(InfoMap) 号码出错
+
+	ENotFound ErrNo = 1004 //@TakeCodeToMap(InfoMap) 找不到
+)
+
+```
+
+生成的代码
+```go
+package demo
+
+var InfoMap = map[int]bool{1003: true, 1004: true}
+
+```

codemsg/codemsg.go

@@ -53,6 +53,9 @@ func genCodeMsg(c *CodeMsg) {
 
 	for _, typeName := range types {
 		g.generateCodeMsg(c, typeName)
+		for mapName, m := range c.SaveCodeToMap {
+			saveTakeFileMap(dir, g.pkg.name, typeName, mapName, m)
+		}
 	}
 
 	saveFile(c, dir, g.format(), "_codemsg.go", types[0])
@@ -62,29 +65,28 @@ func genCodeMsg(c *CodeMsg) {
 }
 
 func saveFile(c *CodeMsg, dir string, src []byte, suffix string, types string) {
-
 	outputName := c.Output
 	if outputName == "" {
 		baseName := fmt.Sprintf("%s%s", types, suffix)
 		outputName = filepath.Join(dir, strings.ToLower(baseName))
 	}
 
-	err := ioutil.WriteFile(outputName, src, 0644)
+	err := ioutil.WriteFile(outputName, src, 0o644)
 	if err != nil {
 		log.Fatalf("writing output: %s", err)
 	}
 }
 
 // generate produces the String method for the named type.
 func (g *Generator) generateCodeMsg(c *CodeMsg, typeName string) {
-
 	values := make([]Value, 0, 100)
 
 	for _, file := range g.pkg.files {
 
 		// Set the state for this run of the walker.
 		file.typeName = typeName
 		file.values = nil
+		file.c = c
 		if file.file != nil {
 			ast.Inspect(file.file, file.genDecl)
 			values = append(values, file.values...)
@@ -100,7 +102,7 @@ func (g *Generator) generateCodeMsg(c *CodeMsg, typeName string) {
 	tmpl.Args = strings.Join(os.Args[2:], " ")
 	tmpl.PkgName = g.pkg.name
 
-	//tmpl.Gen(os.Stdout)
+	// tmpl.Gen(os.Stdout)
 	if err := tmpl.Gen(&g.buf); err != nil {
 		io.Copy(os.Stdout, bytes.NewReader(g.buf.Bytes()))
 	}

codemsg/stringer.go

@@ -95,15 +95,16 @@ type CodeMsg struct {
 	Grpc              bool   `clop:"long" usage:"Generate grpc error type"`
 	StringMethod      string `clop:"long" usage:"String function name" default:"String"`
 	String            bool   `clop:"long" usage:"Generate string function"`
-	CodeMsgStructName string `clop:"long" usage:"turn on the ability to generate codemsg code" default:"CodeMsg"` //TODO
+	CodeMsgStructName string `clop:"long" usage:"turn on the ability to generate codemsg code" default:"CodeMsg"` // TODO
 
-	CodeName string   `clop:"long" usage:"set new code name" default:"Code"`
-	MsgName  string   `clop:"long" usage:"set new message name" default:"Message"`
-	Args     []string `clop:"args=file" usage:"file or dir" default:"["."]"`
+	CodeName      string   `clop:"long" usage:"set new code name" default:"Code"`
+	MsgName       string   `clop:"long" usage:"set new message name" default:"Message"`
+	TakeCodeToMap bool     `clop:"long" usage:"Enable the take to map feature"`
+	Args          []string `clop:"args=file" usage:"file or dir" default:"["."]"`
+	SaveCodeToMap map[string]map[int]bool
 }
 
 func genString(c *CodeMsg) {
-
 	log.SetFlags(0)
 	log.SetPrefix("h2o codemsg: ")
 
@@ -187,6 +188,7 @@ type File struct {
 
 	trimPrefix  string
 	lineComment bool
+	c           *CodeMsg
 }
 
 type Package struct {
@@ -239,14 +241,17 @@ func (g *Generator) addPackage(pkg *packages.Package) {
 
 // generate produces the String method for the named type.
 func (g *Generator) generate(typeName string, c *CodeMsg) {
-
 	values := make([]Value, 0, 100)
 
 	for _, file := range g.pkg.files {
 
 		// Set the state for this run of the walker.
 		file.typeName = typeName
 		file.values = nil
+		file.c = c
+		if file.c == nil {
+			panic("codemsg is nil")
+		}
 		if file.file != nil {
 			ast.Inspect(file.file, file.genDecl)
 			values = append(values, file.values...)
@@ -450,82 +455,34 @@ func (f *File) genDecl(node ast.Node) bool {
 			} else {
 				v.Name = strings.TrimPrefix(v.OriginalName, f.trimPrefix)
 			}
-			f.values = append(f.values, v)
-		}
-	}
-	return false
-}
-
-// Helpers
-
-// usize returns the number of bits of the smallest unsigned integer
-// type that will hold n. Used to create the smallest possible slice of
-// integers to use as indexes into the concatenated strings.
-func usize(n int) int {
-	switch {
-	case n < 1<<8:
-		return 8
-	case n < 1<<16:
-		return 16
-	default:
-		// 2^32 is enough constants for anyone.
-		return 32
-	}
-}
-
-// declareIndexAndNameVars declares the index slices and concatenated names
-// strings representing the runs of values.
-func (g *Generator) declareIndexAndNameVars(runs [][]Value, typeName string) {
-	var indexes, names []string
-	for i, run := range runs {
-		index, name := g.createIndexAndNameDecl(run, typeName, fmt.Sprintf("_%d", i))
-		if len(run) != 1 {
-			indexes = append(indexes, index)
-		}
-		names = append(names, name)
-	}
-	g.Printf("const (\n")
-	for _, name := range names {
-		g.Printf("\t%s\n", name)
-	}
-	g.Printf(")\n\n")
 
-	if len(indexes) > 0 {
-		g.Printf("var (")
-		for _, index := range indexes {
-			g.Printf("\t%s\n", index)
-		}
-		g.Printf(")\n\n")
-	}
-}
+			if f.c == nil {
+				panic("f.c is nil")
+			}
 
-// declareIndexAndNameVar is the single-run version of declareIndexAndNameVars
-func (g *Generator) declareIndexAndNameVar(run []Value, typeName string) {
-	index, name := g.createIndexAndNameDecl(run, typeName, "")
-	g.Printf("const %s\n", name)
-	g.Printf("var %s\n", index)
-}
+			if f.c.TakeCodeToMap {
+				mapName := ""
+				var err error
+				v.Name, mapName, err = parseTakeCodeToMap(v.Name)
+				if err != nil {
+					panic(err)
+				}
+				if f.c.SaveCodeToMap == nil {
+					f.c.SaveCodeToMap = make(map[string]map[int]bool)
+				}
+
+				saveMap := f.c.SaveCodeToMap[mapName]
+				if saveMap == nil {
+					saveMap = make(map[int]bool)
+				}
+				saveMap[int(v.value)] = true
+				f.c.SaveCodeToMap[mapName] = saveMap
+			}
 
-// createIndexAndNameDecl returns the pair of declarations for the run. The caller will add "const" and "var".
-func (g *Generator) createIndexAndNameDecl(run []Value, typeName string, suffix string) (string, string) {
-	b := new(bytes.Buffer)
-	indexes := make([]int, len(run))
-	for i := range run {
-		b.WriteString(run[i].Name)
-		indexes[i] = b.Len()
-	}
-	nameConst := fmt.Sprintf("_%s_name%s = %q", typeName, suffix, b.String())
-	nameLen := b.Len()
-	b.Reset()
-	fmt.Fprintf(b, "_%s_index%s = [...]uint%d{0, ", typeName, suffix, usize(nameLen))
-	for i, v := range indexes {
-		if i > 0 {
-			fmt.Fprintf(b, ", ")
+			f.values = append(f.values, v)
 		}
-		fmt.Fprintf(b, "%d", v)
 	}
-	fmt.Fprintf(b, "}")
-	return b.String(), nameConst
+	return false
 }
 
 // declareNameVars declares the concatenated names string representing all the values in the runs.

codemsg/take_to_map.go

@@ -0,0 +1,73 @@
+package codemsg
+
+import (
+	"bytes"
+	"fmt"
+	"go/format"
+	"os"
+	"path/filepath"
+	"strings"
+	"text/template"
+)
+
+const (
+	takeToMapStart = "@TakeCodeToMap("
+	takeToMapEnd   = ")"
+)
+
+// text样子是 @take_to_map(map名字)
+func parseTakeCodeToMap(text string) (newText, mapName string, err error) {
+	start := strings.Index(text, takeToMapStart)
+	if start == -1 {
+		return text, "", nil
+	}
+
+	end := strings.Index(text[start+len(takeToMapStart):], takeToMapEnd)
+	if end == -1 {
+		return text, "", nil
+	}
+
+	mapName = text[start+len(takeToMapStart) : start+len(takeToMapStart)+end]
+
+	old := text[start : start+len(takeToMapStart)+end+1]
+	newText = strings.ReplaceAll(text, old, "")
+	return newText, mapName, nil
+}
+
+func saveTakeFileMap(dir, packageName, types string, mapName string, m map[int]bool) {
+	baseName := fmt.Sprintf("%s_take_to_map_%s.go", types, mapName)
+	fileName := filepath.Join(dir, strings.ToLower(baseName))
+
+	const tmpl = `package {{.PackageName}}
+		var {{.MapName}} = {{printf "%#v" .MyMap}}`
+	tmplParsed, err := template.New("example").Parse(tmpl)
+	if err != nil {
+		panic(err)
+	}
+
+	// 渲染模板到一个缓冲区
+	var buf bytes.Buffer
+
+	data := struct {
+		PackageName string
+		MapName     string
+		MyMap       map[int]bool
+	}{
+		MapName:     mapName,
+		MyMap:       m,
+		PackageName: packageName,
+	}
+	err = tmplParsed.Execute(&buf, data)
+	if err != nil {
+		panic(err)
+	}
+	sourceCode, err := format.Source(buf.Bytes())
+	if err != nil {
+		panic(err)
+	}
+
+	err = os.WriteFile(fileName, sourceCode, 0o644)
+	if err != nil {
+		panic(err)
+	}
+}