go语言实现腾讯股票获取示例(并发)

package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/text/encoding/simplifiedchinese"
)

// StockData 股票数据结构
type StockData struct {
	Code          string  // 股票代码
	Name          string  // 股票名称
	Price         float64 // 当前价
	Change        float64 // 涨跌额
	ChangePercent float64 // 涨跌幅
	Volume        int64   // 成交量
	Amount        float64 // 成交额
	High          float64 // 最高价
	Low           float64 // 最低价
	Open          float64 // 开盘价
	PreClose      float64 // 前收盘价
}

// 获取股票代码列表
func getStockCodes() ([]string, error) {
	// 模拟股票代码列表
	var codes []string
	// 添加一些上海市场的股票代码
	for i := range 10000 {
		code := fmt.Sprintf("sh%06d", 600000+i)
		codes = append(codes, code)
	}
	// 添加一些深圳市场的股票代码
	for i := range 999 {
		code := fmt.Sprintf("sz%06d", 1+i)
		codes = append(codes, code)
	}
	return codes, nil
}

// 获取单个股票数据
func fetchStockData(code string) (*StockData, error) {
	url := fmt.Sprintf("http://qt.gtimg.cn/q=%s", code)
	log.Printf("开始获取股票 %s 数据,URL: %s", code, url)

	// 创建HTTP客户端
	client := &http.Client{
		Timeout: 5 * time.Second,
	}

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %w", err)
	}

	// 设置适当的请求头
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64)")
	req.Header.Set("Accept", "*/*")

	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("发送请求失败: %w", err)
	}
	defer resp.Body.Close()

	// 检查响应状态码
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("响应状态码异常: %d", resp.StatusCode)
	}

	// 读取响应体
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应体失败: %w", err)
	}

	// 检查响应体大小
	if len(body) == 0 {
		return nil, fmt.Errorf("响应体为空")
	}

	// 将GB2312/GBK编码的响应内容转换为UTF-8
	var content string
	dataUTF8, err := simplifiedchinese.GBK.NewDecoder().Bytes(body)
	if err != nil {
		// 如果转换失败,使用原始内容
		content = string(body)
		log.Printf("[警告] 股票 %s 编码转换失败: %v,将使用原始编码尝试处理", code, err)
	} else {
		content = string(dataUTF8)
	}

	// 检查数据有效性
	if !strings.Contains(content, "~") {
		return nil, fmt.Errorf("数据格式异常: 缺少分隔符~")
	}

	// 提取数据部分 - 优化提取逻辑
	start := strings.Index(content, "=\"")
	end := strings.LastIndex(content, "\"")
	if start == -1 || end == -1 {
		return nil, fmt.Errorf("无法提取有效数据: 找不到引号包裹的数据部分")
	}

	// 提取数据部分
	data := content[start+2 : end]
	if len(data) == 0 {
		return nil, fmt.Errorf("提取的数据为空")
	}

	return parseStockData(code, data)
}

// 解析股票数据
func parseStockData(code, data string) (*StockData, error) {
	log.Printf("开始解析股票 %s 数据", code)

	fields := strings.Split(data, "~")
	if len(fields) < 40 {
		return nil, fmt.Errorf("数据字段不足,期望至少40个字段,实际获取到%d个字段", len(fields))
	}

	stockName := fields[1]
	if stockName == "" {
		log.Printf("[警告] 股票 %s 名称为空,可能是无效股票", code)
	}

	stock := &StockData{
		Code: code,
		Name: stockName,
	}

	// 定义字段映射关系,便于维护
	type fieldParse struct {
		index       int
		fieldPtr    any
		description string
	}

	fieldParses := []fieldParse{
		{3, &stock.Price, "当前价"},
		{31, &stock.Change, "涨跌额"},
		{36, &stock.Volume, "成交量"},
		{37, &stock.Amount, "成交额"},
		{33, &stock.High, "最高价"},
		{34, &stock.Low, "最低价"},
		{5, &stock.Open, "开盘价"},
		{4, &stock.PreClose, "前收盘价"},
	}

	// 解析常规数值字段
	for _, fp := range fieldParses {
		if fp.index >= len(fields) {
			log.Printf("[警告] 股票 %s 数据字段不完整:%s 字段不存在(索引:%d),当前数据共有%d个字段",
				code, fp.description, fp.index, len(fields))
			continue
		}

		fieldValue := fields[fp.index]
		if fieldValue == "" {
			log.Printf("[警告] 股票 %s %s 字段为空值,将使用默认值", code, fp.description)
			continue
		}

		switch ptr := fp.fieldPtr.(type) {
		case *float64:
			val, err := strconv.ParseFloat(fieldValue, 64)
			if err != nil {
				log.Printf("[警告] 股票 %s %s 解析失败: %v,原始值: '%s',将使用默认值",
					code, fp.description, err, fieldValue)
			} else {
				*ptr = val
			}
		case *int64:
			val, err := strconv.ParseInt(fieldValue, 10, 64)
			if err != nil {
				log.Printf("[警告] 股票 %s %s 解析失败: %v,原始值: '%s',将使用默认值",
					code, fp.description, err, fieldValue)
			} else {
				*ptr = val
			}
		}
	}

	// 特殊处理涨跌幅字段
	if len(fields) > 5 {
		changePercentStr := fields[32]
		if changePercentStr != "" {
			// 移除百分号
			changePercentStr = strings.TrimSuffix(changePercentStr, "%")
			val, err := strconv.ParseFloat(changePercentStr, 64)
			if err != nil {
				log.Printf("[警告] 股票 %s 涨跌幅解析失败: %v,原始值: '%s',将使用默认值0",
					code, err, changePercentStr)
			} else {
				stock.ChangePercent = val
			}
		} else {
			log.Printf("[警告] 股票 %s 涨跌幅字段为空,将使用默认值0", code)
		}
	}

	log.Printf("股票 %s 数据解析完成,名称: %s,价格: %.2f", code, stock.Name, stock.Price)
	return stock, nil
}

// 处理股票数据(带重试)
func processStockData(code string, maxRetries int, retryDelay time.Duration) (*StockData, error) {
	var lastErr error

	for i := range maxRetries {
		log.Printf("[尝试 %d/%d] 获取股票 %s 数据", i+1, maxRetries, code)

		// 尝试从API获取数据
		stock, err := fetchStockData(code)
		if err != nil {
			// 记录错误但继续重试
			lastErr = fmt.Errorf("获取股票 %s 数据出错: %w", code, err)
			log.Printf("[重试] 第%d次尝试失败: %v", i+1, lastErr)
		} else if stock != nil && stock.Name != "" && len(stock.Name) > 1 {
			log.Printf("[成功] 第%d次尝试成功获取股票 %s 数据", i+1, code)
			return stock, nil
		} else {
			lastErr = fmt.Errorf("获取股票 %s 数据不完整或无效", code)
			log.Printf("[重试] 第%d次尝试获取到不完整数据", i+1)
		}

		// 如果不是最后一次尝试,则休眠后继续
		if i < maxRetries-1 {
			log.Printf("[重试] 等待 %.2f 秒后进行第%d次尝试", retryDelay.Seconds(), i+2)
			time.Sleep(retryDelay)
		}
	}

	// 所有重试都失败了
	if lastErr != nil {
		return nil, fmt.Errorf("获取股票 %s 数据失败(已重试%d次): %w", code, maxRetries, lastErr)
	}
	return nil, fmt.Errorf("获取股票 %s 数据失败(已重试%d次)", code, maxRetries)
}

// 保存数据到CSV文件
func saveToCSV(stocks []*StockData, filename string) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	// 写入UTF-8 BOM,确保Excel能正确识别编码
	_, err = file.WriteString("\xEF\xBB\xBF")
	if err != nil {
		return err
	}

	// 创建CSV写入器
	writer := csv.NewWriter(file)
	writer.Comma = ','
	writer.UseCRLF = true

	// 写入表头
	headers := []string{"代码", "名称", "当前价", "涨跌额", "涨跌幅", "成交量", "成交额", "最高价", "最低价", "开盘价", "前收盘价"}
	if err := writer.Write(headers); err != nil {
		return err
	}

	// 写入数据行
	for _, stock := range stocks {
		row := []string{
			stock.Code,
			stock.Name, // 已修复的股票名称
			fmt.Sprintf("%.2f", stock.Price),
			fmt.Sprintf("%.2f", stock.Change),
			fmt.Sprintf("%.2f%%", stock.ChangePercent),
			fmt.Sprintf("%d", stock.Volume),
			fmt.Sprintf("%.2f", stock.Amount),
			fmt.Sprintf("%.2f", stock.High),
			fmt.Sprintf("%.2f", stock.Low),
			fmt.Sprintf("%.2f", stock.Open),
			fmt.Sprintf("%.2f", stock.PreClose),
		}
		if err := writer.Write(row); err != nil {
			return err
		}
	}

	writer.Flush()
	return writer.Error()
}

// 显示股票数据
func displayStocks(stocks []*StockData) {
	fmt.Println("\n=== 股票数据统计 ===")
	fmt.Printf("获取到 %d 只股票数据\n\n", len(stocks))

	// 优化列宽,确保中文显示正常
	fmt.Printf("%-10s %-12s %-10s %-10s %-8s %-10s\n",
		"代码", "名称", "当前价", "涨跌额", "涨跌幅", "成交额")
	fmt.Println("----------------------------------------------------------------------")

	// 只显示前5只股票
	limit := 5
	if len(stocks) < limit {
		limit = len(stocks)
	}

	for i := 0; i < limit; i++ {
		stock := stocks[i]
		// 确保股票名称不会溢出列宽
		displayName := stock.Name
		if len([]rune(displayName)) > 6 {
			displayName = string([]rune(displayName)[:6]) + "..."
		}
		// 修复涨跌幅显示格式,确保百分号直接跟在数字后面,没有空格
		fmt.Printf("%-10s %-12s %-10.2f %-10.2f %5.2f%% %-10.2f\n",
			stock.Code, displayName, stock.Price, stock.Change, stock.ChangePercent, stock.Amount)
	}
	fmt.Println()
}

// 验证并规范化命令行参数
func validateParams(maxCount, concurrency *int) {
	// 验证并发数
	if *concurrency <= 0 {
		log.Printf("[配置警告] 并发数设置为 %d,无效(必须大于0),已自动调整为默认值 5", *concurrency)
		*concurrency = 5
	} else if *concurrency > 20 {
		log.Printf("[配置警告] 并发数设置为 %d,超出安全范围(最大20),已自动调整为最大值 20", *concurrency)
		*concurrency = 20
	} else {
		log.Printf("[配置信息] 并发数设置为 %d,在安全范围内", *concurrency)
	}

	// 验证股票数量
	if *maxCount <= 0 {
		log.Printf("[配置警告] 获取股票数量设置为 %d,无效(必须大于0),已自动调整为默认值 10", *maxCount)
		*maxCount = 10
	} else if *maxCount > 100 {
		log.Printf("[配置警告] 获取股票数量设置为 %d,超出推荐范围(最大100),已自动调整为最大值 100", *maxCount)
		*maxCount = 100
	} else {
		log.Printf("[配置信息] 获取股票数量设置为 %d,在推荐范围内", *maxCount)
	}

	log.Printf("参数验证完成: 获取股票数量=%d, 并发数=%d", *maxCount, *concurrency)
}

func main() {
	// 定义命令行参数
	maxCount := flag.Int("n", 10, "获取股票数量")
	concurrency := flag.Int("c", 5, "并发数")
	flag.Parse()

	// 参数验证
	validateParams(maxCount, concurrency)

	fmt.Println("开始获取A股数据...")
	startTime := time.Now()

	// 获取股票代码
	codes, err := getStockCodes()
	if err != nil {
		log.Fatalf("获取股票代码失败: %v", err)
	}

	fmt.Printf("过滤后共有 %d 个股票代码\n", len(codes))
	fmt.Printf("限制获取 %d 只股票数据\n", *maxCount)

	// 限制数量
	if len(codes) > *maxCount {
		codes = codes[:*maxCount]
	}

	fmt.Printf("开始批量获取数据(并发数: %d)...\n", *concurrency)

	// 创建worker池处理股票数据获取
	type stockTask struct {
		code string
	}

	type stockResult struct {
		stock *StockData
		err   error
	}

	// 创建任务通道和结果通道
	taskChan := make(chan stockTask, len(codes))
	resultChan := make(chan stockResult, len(codes))

	// 启动worker池
	var wg sync.WaitGroup
	for i := 0; i < *concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for task := range taskChan {
				stock, err := processStockData(task.code, 3, 300*time.Millisecond)
				resultChan <- stockResult{stock: stock, err: err}
			}
		}()
	}

	// 提交任务
	for _, code := range codes {
		taskChan <- stockTask{code: code}
	}
	close(taskChan)

	// 等待所有worker完成并关闭结果通道
	go func() {
		wg.Wait()
		close(resultChan)
	}()

	// 收集数据和错误
	var validStocks []*StockData
	for result := range resultChan {
		if result.err != nil {
			// 当有错误时,从task中获取code信息
			log.Printf("错误: %v", result.err)
		} else if result.stock != nil {
			validStocks = append(validStocks, result.stock)
		}
	}

	duration := time.Since(startTime)
	fmt.Printf("数据获取完成,耗时: %v\n", duration)

	// 显示数据
	displayStocks(validStocks)

	// 保存到CSV文件
	filename := fmt.Sprintf("stock_data_%s.csv", time.Now().Format("20060102_150405"))
	if err := saveToCSV(validStocks, filename); err != nil {
		log.Fatalf("保存CSV文件失败: %v", err)
	}

	fmt.Printf("数据已保存到: %s\n", filename)
	fmt.Println("提示:CSV文件已使用UTF-8编码(带BOM)保存,可以直接在Excel中打开查看中文内容。")
}
posted @ 2025-11-04 23:13  卓能文  阅读(1)  评论(0)    收藏  举报