runsisi's

technical notes

Go 读写 csv 文档

2019-05-25 runsisigolang

CSV 是一种非常简单的格式,如下所示:

ipv;table;command;chain;proto;src;dst;sport;dport;match;jump;desc
4,6;filter;A;zstor;tcp;;;;6789;;ACCEPT;MON
4,6;filter;A;zstor;tcp;;;;6800:7300;;ACCEPT;OSD,MGR,MDS
4,6;filter;A;zstor;tcp;;;;3300;;ACCEPT;MON >= Nautilus

但 golang 自带的 encoding/csv 包只是简单对字符串解析进行了一些封装,因此用起来算不上有多方便。

第三方的 gocsv 库 虽然实现的并不怎么样,例如没有对外暴露 Encoder/Decoder,没有暴露错误码定义,但用起来还算方便。

在使用之前需要进行必要的设置(不是很理解为何不暴露 Encoder/Decoder):

gocsv.SetCSVReader(func(in io.Reader) gocsv.CSVReader {
		reader := csv.NewReader(in)
		reader.Comma = ';'
		return reader
	})
	gocsv.SetCSVWriter(func(out io.Writer) *gocsv.SafeCSVWriter {
		writer := csv.NewWriter(out)
		writer.Comma = ';'
		return gocsv.NewSafeCSVWriter(writer)
	})

剩下的就非常简单了,首先定义结构体的各字段(注意结构体的 tag 需要与 CSV 文件的表头对应):

type RawRule struct {
	Ipv     string `csv:"ipv"`
	Table   string `csv:"table"`
	Command string `csv:"command"`
	Chain   string `csv:"chain"`
	Proto   string `csv:"proto"`
	Src     string `csv:"src"`
	Dst     string `csv:"dst"`
	Sport   string `csv:"sport"`
	Dport   string `csv:"dport"`
	Match   string `csv:"match"`
	Jump    string `csv:"jump"`
	Desc    string `csv:"desc"`
}

然后直接 unmarshal 操作即可:

var rules []RawRule
if err := gocsv.UnmarshalString(data, &rules); err != nil {
  return RawRules{}, errors.Wrap(err, "failed to parse rules")
}

而 marshal 也非常简单:

buffer := bytes.NewBuffer(nil)
if err := gocsv.Marshal(rules, buffer); err != nil {
  return nil, errors.Wrapf(err, "failed to encode rules")
}