用Golang做点自动化的东西

用protobuf的时候就已经觉得挺好玩的，一个.proto文件，用一个命令加一个language type的参数就能生成相应语言的pb文件，神奇。这阵子闲了一点，调查了下，发现golang原生支持这种东西。核心，go generate。

go generate

简介

go generate命令是go 1.4版本里面新添加的一个命令，当运行go generate时，它将扫描当前目录下的go文件，找出所有包含"//go:generate"的特殊注释，提取并执行该注释后面的命令，命令为可执行程序。

需要看Go的官方使用方法，在命令行下

 $ go help generate
usage: go generate [-run regexp] [-n] [-v] [-x] [build flags] [file.go... | packages]

Generate runs commands described by directives within existing
files. Those commands can run any process but the intent is to
create or update Go source files.

Go generate is never run automatically by go build, go get, go test,
and so on. It must be run explicitly.

....

To convey to humans and machine tools that code is generated,
generated source should have a line that matches the following
regular expression (in Go syntax):

	^// Code generated .* DO NOT EDIT\.$

The line may appear anywhere in the file, but is typically
placed near the beginning so it is easy to find.

....
复制代码

go generatego generate

还有，为了让人类和机器的工具都知道代码是生成，最好在比较容易发现的地方加上

// Code generated .* DO NOT EDIT\.$
复制代码

在我看来，其实go generate就是运行命令生一个文件而已，具体生成的文件是什么格式，有什么用，有什么内容，这都是由开发者自定义的。只不过，最好就是只用来生成和更新go文件，并且在文件内容里面加一个注释来标志这个go文件是自动生成的，仅此而已。

小试

为了验证上面说的话，我决定不走寻常路，不像常规的生成代码了，用go generate来一张二维码。

写一个命令文件

写一个gen.go

package main
//go:generate gen_png

复制代码

嗯，是的，你没看错，包含最后一个空行，只需三行代码

生成二维码

当前目录情况

$ pwd
/Users/cltx/workspace/go/src/github.com/LSivan/go_codes/test/test_generate/gen_png

$ ls
gen.go  main.go

复制代码

go generate

$ go build

$ go generate                                                                                                                                                                                                                                                        1 ↵
gen.go:2: running "gen_png": exec: "gen_png": executable file not found in $PATH
复制代码

执行文件得在PATH目录下，把它移动到GOPATH，确保GOPATH加进了PATH下

$ sudo mv gen_png $GOPATH/bin/
复制代码

再次generate，就能看到目录下多了个二维码

$ go generate  

$ ls
gen.go  hello.png  main.go
复制代码

这里就不贴图啦

正片

我们来写一个json格式的struct生成器。具体来说就是给定一个json，然后根据这个json，生成相应的Go文件。为什么写这个？比如说前后端定义了json的接口，或者接口有所改动，这个东西就很好用了。

json 2 structure

需求大概是这样子，我和前端小明定了一个接口，获取用户风险信息，接口协议如下：

{
  "risk_query_response": {
    "code": "10000",
    "msg": "Success",
    "risk_result": {
      "merchant_fraud": "has_risk" ,
      "merchant_general":"rank_1"
    },
    "risk_result_desc": "{\"has_risk\":\"有风险\",\"rank_1\":\"等级1\",\"71.5\":\"评分71.5\"}"
  },
  "sign": "ERITJKEIJKJHKKKKKKKHJEREEEEEEEEEEE",
  "risk_rule": [
    101,
    1001
  ],
  "time": 1553481619,
  "is_black_user": false
}
复制代码

我这边需要在业务封装一个strcut，然后转json返回给前端，如下

type Risk struct {
	RiskQueryResponse struct {
		Code       string `json:"code"`
		Msg        string `json:"msg"`
		RiskResult struct {
			MerchantFraud   string `json:"merchant_fraud"`
			MerchantGeneral string `json:"merchant_general"`
		} `json:"risk_result"`
		RiskResultDesc string `json:"risk_result_desc"`
	} `json:"risk_query_response"`
	Sign        string `json:"sign"`
	RiskRule    []int  `json:"risk_rule"`
	Time        int    `json:"time"`
	IsBlackUser bool   `json:"is_black_user"`
}
复制代码

手写当然简单，不到三分钟就写好了。但是几十上百个接口的时候，这就可怕了。为此，我们来试下，go generate。

编写命令(可执行)文件

这个命令(可执行文件)的本质其实就是解析json，然后得到一个Go文件。

先是读取文件，解析json

if f == "" {
    panic("file can not be nil")
}

jsonFile, err := os.Open(f)
if err != nil {
    panic(fmt.Sprintf("open file error:%s", err.Error()))
}
fi, err := jsonFile.Stat()
if err != nil {
    panic(fmt.Sprintf("get file stat error:%s",err.Error()))
}

if fi.Size() > 40960 {
    panic("json too big")
}
data := make([]byte, 40960)
bio := bufio.NewReader(jsonFile)
n, err := bio.Read(data)
if err != nil {
    panic(err)
}
fmt.Println(string(data[:n]))

m := new(map[string]interface{})
err = json.Unmarshal(data[:n], m)
if err != nil {
    panic(fmt.Sprintf("unmarshal json error:%s", err.Error()))
}
复制代码

反射得到json各对键值的类型，这里有个细节，go会用float64来接收数值类型；此外，这里用到了一个转驼峰命名的第三方库strcase

field := ""
for k, v := range *m {
    t := reflect.TypeOf(v)
    kind := t.Kind()
    fieldType := t.String()
    switch kind {
    case reflect.String:
        field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
    case reflect.Map:
        field += strcase.ToCamel(k) + " struct {\n"
        fields := parserMap(v.(map[string]interface{}))
        field += fields
        field += "}	" +fmt.Sprintf("`json:\"%s\"`\n",k)
    case reflect.Slice:
        fieldType = parserSlice(v.([]interface{}))
        field += strcase.ToCamel(k) + "[]" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
    case reflect.Float64:
        fieldType = "int"
        field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
    case reflect.Bool:
        field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
    default:
        fmt.Println("other kind", k, kind)
    }
}

复制代码

对map和slice类型做特殊处理，让他们变成相应的struct

func parserMap(m map[string]interface{}) string {
	field := ""
	for k,v := range m {
		t := reflect.TypeOf(v)
		kind := t.Kind()
		fieldType := t.String()
		switch kind {
		case reflect.String:
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Map:
			field += strcase.ToCamel(k) + " struct {\n"
			fields := parserMap(v.(map[string]interface{}))
			field += fields
			field += "}	" +fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Slice:
			parserSlice(v.([]interface{}))
		case reflect.Float64:
			fieldType = "int"
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Bool:
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		default:
			fmt.Println("other kind", k, kind)
		}
	}
	return field
}

func parserSlice(s []interface{}) string {
	field := ""
	for k,v := range s {
		t := reflect.TypeOf(v)
		kind := t.Kind()
		fieldType := t.String()
		switch kind {
		case reflect.String:
			return fieldType
		case reflect.Float64:
			fieldType = "int"
			return fieldType
		case reflect.Bool:
			return fieldType
		default:
			fmt.Println("other kind", k, kind)
		}
	}
	return field
}
复制代码

写入到go文件中，同时别忘了遵循下官方的约束，在文件开头加个自动生成声明

fileName := strings.Split(fi.Name(), ".")[0]
goFile := fmt.Sprintf(output+"/%s.go", fileName)
f, _ := os.Create(goFile)

template := "// Code generated json_2_struct. DO NOT EDIT.\n" +
	"package %s\n" +
	"\n" +
	" type %s struct {\n"

_, _ = f.Write([]byte(fmt.Sprintf(template+field+"}", pack, strcase.ToCamel(fileName))))
_ = f.Close()
复制代码

最后fmt一下生成的go文件

cmd := exec.Command("go", "fmt", goFile)
err = cmd.Run()
if err != nil {
	panic(err)
}
复制代码

编译好之后放到GOPATH

$ go build -o json_2_struct                                                                              
$ sudo mv json_2_struct $GOPATH/bin/ 
复制代码

第一步大功告成

写个脚本用起来

很遗憾的是，go generate命令扫描的必定是go文件，因此脚本得先写个go文件，然后go文件里面增加go:generate的注释，然后执行go generate，脚本如下

#!/bin/bash
echo "package main
//go:generate json_2_struct -file=$1 -output=$2

" > tmp.go

go generate

rm tmp.go
复制代码

P.S. 没有做非法校验

试下效果

让我们使用脚本生成go文件（P.S. 我将脚本移到$PATH下了）

$ export_go_file /Users/cltx/workspace/go/src/github.com/LSivan/go_codes/test/test_generate/data/risk.json /Users/cltx/workspace/go/src/github.com/LSivan/go_codes/test/test_generate/data/                                                
{
  "risk_query_response": {
    "code": "10000",
    "msg": "Success",
    "risk_result": {
      "merchant_fraud": "has_risk" ,
      "merchant_general":"rank_1"
    },
    "risk_result_desc": "{\"has_risk\":\"有风险\",\"rank_1\":\"等级1\",\"71.5\":\"评分71.5\"}"
  },
  "sign": "ERITJKEIJKJHKKKKKKKHJEREEEEEEEEEEE",
  "risk_rule": [
    101,
    1001
  ],
  "time": 1553481619,
  "is_black_user": false
}

$ ls                                                                                                                                                                                                                                      [21:35:04]
risk.go   risk.json

$ cat risk.go                                                                                                                                                                                                                             [21:35:42]
// Code generated test_generate. DO NOT EDIT.
package main

type Risk struct {
        RiskQueryResponse struct {
                Code       string `json:"code"`
                Msg        string `json:"msg"`
                RiskResult struct {
                        MerchantFraud   string `json:"merchant_fraud"`
                        MerchantGeneral string `json:"merchant_general"`
                } `json:"risk_result"`
                RiskResultDesc string `json:"risk_result_desc"`
        } `json:"risk_query_response"`
        Sign        string `json:"sign"`
        RiskRule    []int  `json:"risk_rule"`
        Time        int    `json:"time"`
        IsBlackUser bool   `json:"is_black_user"`
}

复制代码

OK，大功告成

总结

做了一大半之后，发现原来已经早有实现了，尴尬，硬着头皮重复写了个轮子，且当学习吧。

或许会有疑问，为什么不直接写脚本执行那个可执行文件，非得用go generate？而且为什么不用脚本属性更强的python呢？

问得好，于是看了两个generate tool的代码，在stringer的代码里找到了一点痕迹，我认为go generate适用于go file -> go file的情况，因为golang有ast支持。

举个例子，对于hello.go，分别有函数A和函数B，只希望生成函数B的表格单元测试代码，这种情况用golang 原生的ast包 + go:generate就十分方便快捷。

最后附上全代码

package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/iancoleman/strcase"
	"os"
	"os/exec"
	"reflect"
	"strings"
)

var output string
var pack string
var f string

func main() {

	flag.StringVar(&output, "output", "", "output dir")
	flag.StringVar(&pack, "package", "main", "package")
	flag.StringVar(&f, "file", "", "json file")
	flag.Parse()

	if f == "" {
		panic("file can not be nil")
	}

	if output == "" {
		panic("output dir is nil")
	}
	jsonFile, err := os.Open(f)
	if err != nil {
		panic(fmt.Sprintf("open file error:%s", err.Error()))
	}
	fi, err := jsonFile.Stat()
	if err != nil {
		panic(fmt.Sprintf("get file stat error:%s",err.Error()))
	}

	if fi.Size() > 40960 {
		panic("json too big")
	}
	data := make([]byte, 40960)
	bio := bufio.NewReader(jsonFile)
	n, err := bio.Read(data)
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data[:n]))

	m := new(map[string]interface{})
	err = json.Unmarshal(data[:n], m)
	if err != nil {
		panic(fmt.Sprintf("unmarshal json error:%s", err.Error()))
	}

	field := ""
	for k, v := range *m {
		t := reflect.TypeOf(v)
		kind := t.Kind()
		fieldType := t.String()
		switch kind {
		case reflect.String:
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Map:
			field += strcase.ToCamel(k) + " struct {\n"
			fields := parserMap(v.(map[string]interface{}))
			field += fields
			field += "}	" +fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Slice:
			fieldType = parserSlice(v.([]interface{}))
			field += strcase.ToCamel(k) + "[]" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Float64:
			fieldType = "int"
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Bool:
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		default:
			fmt.Println("other kind", k, kind)
		}
	}

	fileName := strings.Split(fi.Name(), ".")[0]
	goFile := fmt.Sprintf(output+"%s.go", fileName)
	file, _ := os.Create(goFile)

	template := "// Code generated test_generate. DO NOT EDIT.\n" +
		"package %s\n" +
		"\n" +
		" type %s struct {\n"

	_, _ = file.Write([]byte(fmt.Sprintf(template+field+"}", pack, strcase.ToCamel(fileName))))
	_ = file.Close()
	
	cmd := exec.Command("go", "fmt", goFile)
	err = cmd.Run()
	if err != nil {
		panic(err)
	}

}


func parserMap(m map[string]interface{}) string {
	field := ""
	for k,v := range m {
		t := reflect.TypeOf(v)
		kind := t.Kind()
		fieldType := t.String()
		switch kind {
		case reflect.String:
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Map:
			field += strcase.ToCamel(k) + " struct {\n"
			fields := parserMap(v.(map[string]interface{}))
			field += fields
			field += "}	" +fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Slice:
			parserSlice(v.([]interface{}))
		case reflect.Float64:
			fieldType = "int"
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		case reflect.Bool:
			field += strcase.ToCamel(k) + "	" + fieldType + fmt.Sprintf("`json:\"%s\"`\n",k)
		default:
			fmt.Println("other kind", k, kind)
		}
	}
	return field
}

func parserSlice(s []interface{}) string {
	field := ""
	for k,v := range s {
		t := reflect.TypeOf(v)
		kind := t.Kind()
		fieldType := t.String()
		switch kind {
		case reflect.String:
			return fieldType
		case reflect.Float64:
			fieldType = "int"
			return fieldType
		case reflect.Bool:
			return fieldType
		default:
			fmt.Println("other kind", k, kind)
		}
	}
	return field
}

复制代码