package tushare import ( "bytes" "encoding/json" "fmt" "io" "net/http" "os" "strings" "time" "github.com/jedib0t/go-pretty/v6/table" ) // TushareClient Tushare API 客户端 type TushareClient struct { Token string `json:"token"` // API 访问令牌 BaseUrl string `json:"base_url"` // API 基础 URL } // TushareReq Tushare API 请求结构 type TushareReq struct { APIName string `json:"api_name"` // API 接口名称 Token string `json:"token"` // API 访问令牌 Params map[string]any `json:"params"` // 请求参数 Fields []string `json:"fields"` // 返回字段列表 } // TushareResp Tushare API 响应结构 type TushareResp struct { RequestID string `json:"request_id"` // 请求 ID Code int `json:"code"` // 响应码 (0 表示成功) Data *TushareRespData `json:"data"` // 响应数据 Msg string `json:"msg"` // 响应消息 } // TushareRespData Tushare API 响应数据结构 type TushareRespData struct { Fields []string `json:"fields"` // 字段列表 Headers []string `json:"headers"` // 表头(中文描述) Items [][]any `json:"items"` // 数据项(二维数组) HasMore bool `json:"has_more"` // 是否有更多数据 Count int `json:"count"` // 返回数据条数 } // NewClient 创建并初始化 Tushare 客户端 // 返回配置好 Token 和基础 URL 的客户端实例 func NewClient(token string) *TushareClient { if token == "" { token = os.Getenv("TUSHARE_TOKEN") } return &TushareClient{ Token: token, BaseUrl: "http://api.tushare.pro", } } // Do 执行 Tushare API 请求 // req: 请求参数,包含 API 名称、参数等 // fieldsVals: 字段配置列表,每个元素是一个 map,key 为字段名,value 为字段中文描述 // 返回:响应数据和错误信息 func (c *TushareClient) Do(req TushareReq, fieldsVals []map[string]string) (*TushareRespData, error) { // 构建请求体 payload := TushareReq{ APIName: req.APIName, Token: c.Token, Params: req.Params, } // 提取字段名和对应的中文表头 var fields []string var headers []string for _, setting := range fieldsVals { for key, value := range setting { fields = append(fields, key) headers = append(headers, value) } } payload.Fields = fields // 序列化请求体为 JSON reqBytes, err := json.Marshal(payload) if err != nil { return nil, fmt.Errorf("#100 请求序列化失败:%w", err) } // 创建 HTTP 请求 client := &http.Client{ Timeout: 30 * time.Second, // 设置 30 秒超时 } resp, err := client.Post(c.BaseUrl, "application/json", bytes.NewReader(reqBytes)) if err != nil { return nil, fmt.Errorf("#100 发送请求失败:%w", err) } defer resp.Body.Close() // 读取响应体 body, err := io.ReadAll(resp.Body) if err != nil { return nil, fmt.Errorf("#101 读取响应失败:%w", err) } // 解析响应 JSON var tushareResp TushareResp if err := json.Unmarshal(body, &tushareResp); err != nil { return nil, fmt.Errorf("#101 响应解析失败:%v - %s", err, string(body)) } // 检查响应码 if tushareResp.Code != 0 { return nil, fmt.Errorf("#102 API 返回错误:%s", tushareResp.Msg) } // 设置表头信息 if tushareResp.Data != nil { tushareResp.Data.Headers = headers } return tushareResp.Data, nil } // Do 执行 Tushare API 请求 // req: 请求参数,包含 API 名称、参数等 // fieldsVals: 字段配置列表,每个元素是一个 map,key 为字段名,value 为字段中文描述 // 返回:响应数据和错误信息 func (c *TushareClient) Exec(ApiName string, fields string, Params map[string]any) (*TushareRespData, error) { fieldsList := strings.Split(fields, ",") if len(fieldsList) == 0 { return nil, fmt.Errorf("#100 字段列表不能为空") } // 构建请求体 payload := TushareReq{ APIName: ApiName, Token: c.Token, Params: Params, Fields: fieldsList, } // 序列化请求体为 JSON reqBytes, err := json.Marshal(payload) if err != nil { return nil, fmt.Errorf("#100 请求序列化失败:%w", err) } // 创建 HTTP 请求 client := &http.Client{ Timeout: 30 * time.Second, // 设置 30 秒超时 } resp, err := client.Post(c.BaseUrl, "application/json", bytes.NewReader(reqBytes)) if err != nil { return nil, fmt.Errorf("#100 发送请求失败:%w", err) } defer resp.Body.Close() // 读取响应体 body, err := io.ReadAll(resp.Body) if err != nil { return nil, fmt.Errorf("#101 读取响应失败:%w", err) } // 解析响应 JSON var tushareResp TushareResp if err := json.Unmarshal(body, &tushareResp); err != nil { return nil, fmt.Errorf("#101 响应解析失败:%v - %s", err, string(body)) } // 检查响应码 if tushareResp.Code != 0 { return nil, fmt.Errorf("#102 API 返回错误:%s", tushareResp.Msg) } // 设置表头信息 return tushareResp.Data, nil } // Map 将响应数据转换为 map 切片 // 每个 map 代表一行数据,key 为字段名,value 为对应的值 func (c *TushareRespData) Map() []map[string]any { result := make([]map[string]any, 0) if c == nil || len(c.Items) == 0 || len(c.Fields) == 0 { return result } for _, item := range c.Items { rowMap := make(map[string]any, len(c.Fields)) for idx, fn := range c.Fields { if idx < len(item) { rowMap[fn] = item[idx] } } result = append(result, rowMap) } return result } // Json 将响应数据转换为 JSON 字节数组 // 返回:JSON 格式的字节数据和错误信息 func (c *TushareRespData) Json() ([]byte, error) { data := c.Map() return json.MarshalIndent(data, "", " ") } // Output 将响应数据格式化为表格输出 // title: 表格标题 // 返回:格式化的表格写入器 func (c *TushareRespData) Output(title string) table.Writer { tw := table.NewWriter() tw.SetStyle(table.StyleLight) tw.SetTitle(title) // 构建表头行 headerRow := make(table.Row, 0, len(c.Headers)) for idx, header := range c.Headers { headerRow = append(headerRow, header+"("+c.Fields[idx]+")") } tw.AppendHeader(headerRow) // 添加数据行 for _, item := range c.Items { tw.AppendRow(item) } return tw }