用Golang做點自動化的東西

liangxingwei發表於2019-03-25

用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 generate就是執行宣告在檔案裡面的命令,這些命令的意圖是生成或更新go檔案(我的理解是:官方希望的約束);go generate不會被類似go get,go build,go test等的命令觸發執行,必須由開發者顯式使用。

還有,為了讓人類和機器的工具都知道程式碼是生成,最好在比較容易發現的地方加上

// Code generated .* DO NOT EDIT\.$
複製程式碼

在我看來,其實go generate就是執行命令生一個檔案而已,具體生成的檔案是什麼格式,有什麼用,有什麼內容,這都是由開發者自定義的。只不過,最好就是隻用來生成和更新go檔案,並且在檔案內容裡面加一個註釋來標誌這個go檔案是自動生成的,僅此而已。

小試

為了驗證上面說的話,我決定不走尋常路,不像常規的生成程式碼了,用go generate來一張二維碼。

寫一個命令檔案

package main

import (
	"github.com/skip2/go-qrcode"
	"os"
)

func main() {
	data,_ := qrcode.Encode("go generate 生成的圖片",qrcode.Medium,256)
	f,_ := os.Create("hello.png")
	f.Write(data)
    f.Close()
}

// P.S. 瘋狂忽略error的demo,真正寫業務的時候請勿模仿 
複製程式碼

以上程式碼很簡單,用第三方庫qrcode生成了一張叫hello.png的二維碼.

寫一個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
}

複製程式碼

相關文章