qsdk/tushare/new.go

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
}

// Exec 执行 Tushare API 请求
// ApiName: API 名称
// fields: 字段列表，用逗号分隔
// Params: 请求参数
// 返回：响应数据和错误信息
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
}