export excel

exportor/data.go

@@ -7,6 +7,7 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/axiaoxin-com/goutils"
 	"github.com/axiaoxin-com/x-stock/model"
 )
 
@@ -76,6 +77,16 @@ type Data struct {
 	EPSPredict string `json:"eps_predict"               csv:"每股收益预测"`
 }
 
+// GetHeaderValueMap 获取以 csv tag 为 key 的 Data map
+func (d Data) GetHeaderValueMap() map[string]interface{} {
+	return goutils.StructToMap(&d, "csv")
+}
+
+// GetHeaders 获取 csv tag 列表
+func (d Data) GetHeaders() []string {
+	return goutils.StructTagList(&d, "csv")
+}
+
 // NewData 创建 Data 对象
 func NewData(stock model.Stock) Data {
 	var rightPrice interface{} = "--"
@@ -155,6 +166,17 @@ func (d DataList) SortByHV() {
 	})
 }
 
+// IndustryList 获取行业分类列表
+func (d DataList) GetIndustryList() []string {
+	result := []string{}
+	for _, stock := range d {
+		if !goutils.IsStrInSlice(stock.Industry, result) {
+			result = append(result, stock.Industry)
+		}
+	}
+	return result
+}
+
 // NewDataList 创建要导出的数据列表
 func NewDataList(ctx context.Context, stocks model.StockList) (result DataList) {
 	for _, s := range stocks {

exportor/excel.go

@@ -0,0 +1,183 @@
+// 导出结果为 excel
+
+package exportor
+
+import (
+	"context"
+	"strings"
+
+	"github.com/360EntSecGroup-Skylar/excelize/v2"
+	"github.com/axiaoxin-com/logging"
+)
+
+var (
+
+	// HeaderStyle 表头样式
+	HeaderStyle = &excelize.Style{
+		Border: []excelize.Border{
+			{Type: "left", Color: "000000", Style: 1},
+			{Type: "right", Color: "000000", Style: 1},
+			{Type: "top", Color: "000000", Style: 1},
+			{Type: "bottom", Color: "000000", Style: 1},
+		},
+		Fill: excelize.Fill{
+			Type:    "pattern",
+			Pattern: 1,
+			Color:   []string{"FFCCCC"},
+			Shading: 0,
+		},
+		Font: &excelize.Font{
+			Bold: true,
+		},
+		Alignment: &excelize.Alignment{
+			Horizontal:      "center",
+			JustifyLastLine: true,
+			Vertical:        "center",
+			WrapText:        true,
+		},
+	}
+	// BodyStyle 表格Style
+	BodyStyle = &excelize.Style{
+		Alignment: &excelize.Alignment{
+			Horizontal:      "left",
+			JustifyLastLine: true,
+			Vertical:        "center",
+			WrapText:        true,
+		},
+	}
+)
+
+// ExportExcel 导出 excel
+func (e Exportor) ExportExcel(ctx context.Context, filename string) (result []byte, err error) {
+	f := excelize.NewFile()
+
+	// 创建全部数据表
+	defaultSheet := "Sheet1"
+	lowPriceSheet := "30元内"
+	hv1Sheet := "波动率低于0.1"
+	hv2Sheet := "波动率0.1-0.5"
+	hv3Sheet := "波动率高于0.5"
+	sheets := []string{defaultSheet, lowPriceSheet, hv1Sheet, hv2Sheet, hv3Sheet}
+	// 添加行业
+	for _, industry := range e.Stocks.GetIndustryList() {
+		sheets = append(sheets, industry+"行业")
+	}
+
+	headers := e.Stocks[0].GetHeaders()
+	headerStyle, err := f.NewStyle(HeaderStyle)
+	if err != nil {
+		logging.Error(ctx, "New HeaderStyle error:"+err.Error())
+	}
+	bodyStyle, err := f.NewStyle(BodyStyle)
+	if err != nil {
+		logging.Error(ctx, "New BodyStyle error:"+err.Error())
+	}
+	// 创建 sheet
+	for _, sheet := range sheets {
+		f.NewSheet(sheet)
+		for i, header := range headers {
+			// 设置列宽
+			colNum := i + 1
+			width := 30.0
+			switch header {
+			case "主营构成", "每股收益预测":
+				width = 45.0
+			case "公司信息":
+				width = 75.0
+			}
+			col, err := excelize.ColumnNumberToName(colNum)
+			if err != nil {
+				logging.Error(ctx, "CoordinatesToCellName error:"+err.Error())
+			}
+			f.SetColWidth(sheet, col, col, width)
+			// 设置表头行高
+			rowNum := 1
+			height := 20.0
+			f.SetRowHeight(sheet, rowNum, height)
+		}
+
+		// 设置表头样式
+		hcell, err := excelize.CoordinatesToCellName(1, 1)
+		if err != nil {
+			logging.Error(ctx, "CoordinatesToCellName error:"+err.Error())
+			continue
+		}
+		vcell, err := excelize.CoordinatesToCellName(len(headers), 1)
+		if err != nil {
+			logging.Error(ctx, "CoordinatesToCellName error:"+err.Error())
+			continue
+		}
+		f.SetCellStyle(sheet, hcell, vcell, headerStyle)
+
+		// 设置表格样式
+		hcell, err = excelize.CoordinatesToCellName(1, 2)
+		if err != nil {
+			logging.Error(ctx, "CoordinatesToCellName error:"+err.Error())
+			continue
+		}
+		vcell, err = excelize.CoordinatesToCellName(len(headers), len(e.Stocks))
+		if err != nil {
+			logging.Error(ctx, "CoordinatesToCellName error:"+err.Error())
+			continue
+		}
+		f.SetCellStyle(sheet, hcell, vcell, bodyStyle)
+	}
+
+	// 写 header
+	for _, sheet := range sheets {
+		for i, header := range headers {
+			axis, err := excelize.CoordinatesToCellName(i+1, 1)
+			if err != nil {
+				logging.Error(ctx, "CoordinatesToCellName error:"+err.Error())
+				continue
+			}
+			f.SetCellValue(sheet, axis, header)
+		}
+	}
+	// 写 body
+	e.Stocks.SortByROE()
+	for _, sheet := range sheets {
+		row := 2
+		for _, stock := range e.Stocks {
+			switch sheet {
+			case defaultSheet:
+			case lowPriceSheet:
+				if stock.Price > 30 {
+					continue
+				}
+			case hv1Sheet:
+				if stock.HV > 0.1 {
+					continue
+				}
+			case hv2Sheet:
+				if stock.HV <= 0.1 || stock.HV > 0.5 {
+					continue
+				}
+			case hv3Sheet:
+				if stock.HV <= 0.5 {
+					continue
+				}
+			}
+			if strings.HasSuffix(sheet, "行业") && !strings.Contains(sheet, stock.Industry) {
+				continue
+			}
+			headerValueMap := stock.GetHeaderValueMap()
+			for k, header := range headers {
+				col := k + 1
+				axis, err := excelize.CoordinatesToCellName(col, row)
+				if err != nil {
+					logging.Error(ctx, "CoordinatesToCellName error:"+err.Error())
+					continue
+				}
+				value := headerValueMap[header]
+				f.SetCellValue(sheet, axis, value)
+			}
+			row++
+		}
+	}
+
+	buf, err := f.WriteToBuffer()
+	result = buf.Bytes()
+	err = f.SaveAs(filename)
+	return
+}

exportor/excel_test.go

@@ -0,0 +1,13 @@
+package exportor
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestExportExcel(t *testing.T) {
+	e := Exportor{}
+	_, err := e.ExportExcel(_ctx, "./test.xlsx")
+	require.Nil(t, err)
+}

exportor/exporter.go

@@ -39,13 +39,17 @@ func Export(ctx context.Context, exportFilename string) {
 	switch fileext {
 	case ".json":
 		exportType = "json"
-	case ".csv":
+	case ".csv", ".txt":
 		exportType = "csv"
+	case ".xlsx", ".xls":
+		exportType = "excel"
 	}
 	if _, err := os.Stat(filedir); os.IsNotExist(err) {
 		os.Mkdir(filedir, 0755)
 	}
+
 	logging.Infof(ctx, "x-stock exportor start export selected stocks to %s", exportFilename)
+	var err error
 	stocks, err := parser.AutoFilterStocks(ctx)
 	if err != nil {
 		logging.Fatal(ctx, err.Error())
@@ -55,15 +59,14 @@ func Export(ctx context.Context, exportFilename string) {
 
 	switch exportType {
 	case "json":
-		_, err := e.ExportJSON(ctx, exportFilename)
-		if err != nil {
-			logging.Fatal(ctx, err.Error())
-		}
+		_, err = e.ExportJSON(ctx, exportFilename)
 	case "csv":
-		_, err := e.ExportCSV(ctx, exportFilename)
-		if err != nil {
-			logging.Fatal(ctx, err.Error())
-		}
+		_, err = e.ExportCSV(ctx, exportFilename)
+	case "excel":
+		_, err = e.ExportExcel(ctx, exportFilename)
+	}
+	if err != nil {
+		logging.Fatal(ctx, err.Error())
 	}
 
 	latency := time.Now().Sub(beginTime).Milliseconds()

go.mod

@@ -3,7 +3,8 @@ module github.com/axiaoxin-com/x-stock
 go 1.16
 
 require (
-	github.com/axiaoxin-com/goutils v1.0.7
+	github.com/360EntSecGroup-Skylar/excelize/v2 v2.4.0
+	github.com/axiaoxin-com/goutils v1.0.10
 	github.com/axiaoxin-com/logging v1.2.9
 	github.com/denisenkom/go-mssqldb v0.10.0 // indirect
 	github.com/gin-gonic/gin v1.7.1 // indirect
@@ -23,10 +24,10 @@ require (
 	github.com/spf13/afero v1.6.0 // indirect
 	github.com/stretchr/testify v1.7.0
 	github.com/ugorji/go v1.2.5 // indirect
+	go.opentelemetry.io/otel v0.20.0 // indirect
 	go.uber.org/zap v1.16.0
 	golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b // indirect
-	golang.org/x/net v0.0.0-20210415231046-e915ea6b2b7d // indirect
-	golang.org/x/sys v0.0.0-20210421221651-33663a62ff08 // indirect
+	golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gorm.io/driver/mysql v1.0.5 // indirect
 	gorm.io/driver/postgres v1.0.8 // indirect