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 client *http.Client } // 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", client: &http.Client{Timeout: 30 * time.Second}, } } func (c *TushareClient) httpDo() *http.Client { if c.client != nil { return c.client } return &http.Client{Timeout: 30 * time.Second} } // postTushare 发送 JSON 请求并解析为 TushareResp(不含业务码以外的处理)。 func (c *TushareClient) postTushare(reqBytes []byte) (*TushareResp, error) { resp, err := c.httpDo().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) } 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, nil } // 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, } // 提取字段名和对应的中文表头 n := 0 for _, setting := range fieldsVals { n += len(setting) } fields := make([]string, 0, n) headers := make([]string, 0, n) for _, setting := range fieldsVals { for key, value := range setting { fields = append(fields, key) headers = append(headers, value) } } payload.Fields = fields reqBytes, err := json.Marshal(payload) if err != nil { return nil, fmt.Errorf("#100 请求序列化失败:%w", err) } tushareResp, err := c.postTushare(reqBytes) if err != nil { return nil, err } if tushareResp.Data != nil { tushareResp.Data.Headers = headers } return tushareResp.Data, nil } // Exec 执行 Tushare API 请求 // ApiName: API 名称 // fields: 字段列表,用逗号分隔 // Params: 请求参数 // 返回:响应数据和错误信息 func (c *TushareClient) Exec(ApiName string, fields string, Params map[string]any) (*TushareRespData, error) { fields = strings.TrimSpace(fields) if fields == "" { return nil, fmt.Errorf("#100 字段列表不能为空") } raw := strings.Split(fields, ",") fieldsList := make([]string, 0, len(raw)) for _, f := range raw { f = strings.TrimSpace(f) if f != "" { fieldsList = append(fieldsList, f) } } if len(fieldsList) == 0 { return nil, fmt.Errorf("#100 字段列表不能为空") } payload := TushareReq{ APIName: ApiName, Token: c.Token, Params: Params, Fields: fieldsList, } reqBytes, err := json.Marshal(payload) if err != nil { return nil, fmt.Errorf("#100 请求序列化失败:%w", err) } tushareResp, err := c.postTushare(reqBytes) if err != nil { return nil, err } return tushareResp.Data, nil } // Map 将响应数据转换为 map 切片 // 每个 map 代表一行数据,key 为字段名,value 为对应的值 func (c *TushareRespData) Map() []map[string]any { if c == nil || len(c.Items) == 0 || len(c.Fields) == 0 { return make([]map[string]any, 0) } result := make([]map[string]any, 0, len(c.Items)) 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) if c == nil { return tw } n := len(c.Fields) if nh := len(c.Headers); nh < n { n = nh } headerRow := make(table.Row, 0, n) for idx := 0; idx < n; idx++ { headerRow = append(headerRow, c.Headers[idx]+"("+c.Fields[idx]+")") } tw.AppendHeader(headerRow) // 添加数据行 for _, item := range c.Items { tw.AppendRow(item) } return tw }