增加指数代码的判断

codes.go

@@ -2,16 +2,17 @@ package tdx
 
 import (
 	"errors"
-	"github.com/injoyai/conv"
-	"github.com/injoyai/ios/client"
-	"github.com/injoyai/logs"
-	"github.com/injoyai/tdx/protocol"
-	"github.com/robfig/cron/v3"
 	"iter"
 	"math"
 	"os"
 	"path/filepath"
 	"time"
+
+	"github.com/injoyai/conv"
+	"github.com/injoyai/ios/client"
+	"github.com/injoyai/logs"
+	"github.com/injoyai/tdx/protocol"
+	"github.com/robfig/cron/v3"
 	"xorm.io/core"
 	"xorm.io/xorm"
 )
@@ -24,6 +25,8 @@ type ICodes interface {
 	GetStockCodes(limit ...int) []string
 	GetETFs(limit ...int) CodeModels
 	GetETFCodes(limit ...int) []string
+	GetIndexes(limits ...int) CodeModels
+	GetIndexCodes(limits ...int) []string
 }
 
 // DefaultCodes 增加单例,部分数据需要通过Codes里面的信息计算
@@ -134,6 +137,8 @@ func NewCodes(c *Client, db *xorm.Engine) (*Codes, error) {
 	return cc, cc.Update(true)
 }
 
+var _ ICodes = &Codes{}
+
 type Codes struct {
 	*Client                         //客户端
 	db        *xorm.Engine          //数据库实例
@@ -205,6 +210,26 @@ func (this *Codes) GetETFCodes(limits ...int) []string {
 	return this.GetETFs(limits...).Codes()
 }
 
+// GetIndexes 获取基金代码,sz159xxx,sh510xxx,sh511xxx
+func (this *Codes) GetIndexes(limits ...int) CodeModels {
+	limit := conv.Default(-1, limits...)
+	ls := []*CodeModel(nil)
+	for _, m := range this.list {
+		code := m.FullCode()
+		if protocol.IsIndex(code) {
+			ls = append(ls, m)
+		}
+		if limit > 0 && len(ls) >= limit {
+			break
+		}
+	}
+	return ls
+}
+
+func (this *Codes) GetIndexCodes(limits ...int) []string {
+	return this.GetIndexes(limits...).Codes()
+}
+
 func (this *Codes) AddExchange(code string) string {
 	return protocol.AddPrefix(code)
 }

codes_v2.go

@@ -2,6 +2,11 @@ package tdx
 
 import (
 	"errors"
+	"iter"
+	"os"
+	"path/filepath"
+	"time"
+
 	"github.com/injoyai/base/maps"
 	"github.com/injoyai/base/types"
 	"github.com/injoyai/conv"
@@ -12,10 +17,6 @@ import (
 	"github.com/injoyai/tdx/internal/xorms"
 	"github.com/injoyai/tdx/protocol"
 	"github.com/robfig/cron/v3"
-	"iter"
-	"os"
-	"path/filepath"
-	"time"
 	"xorm.io/xorm"
 )
 
@@ -149,12 +150,13 @@ type Codes2 struct {
 		内部字段
 	*/
 
-	c      *Client                           //
-	db     *xorms.Engine                     //
-	stocks types.List[*CodeModel]            //缓存
-	etfs   types.List[*CodeModel]            //缓存
-	all    types.List[*CodeModel]            //缓存
-	m      *maps.Generic[string, *CodeModel] //缓存
+	c       *Client                           //
+	db      *xorms.Engine                     //
+	stocks  types.List[*CodeModel]            //股票缓存
+	etfs    types.List[*CodeModel]            //etf缓存
+	indexes types.List[*CodeModel]            //指数缓存
+	all     types.List[*CodeModel]            //全部缓存
+	m       *maps.Generic[string, *CodeModel] //缓存
 }
 
 func (this *Codes2) Get(code string) *CodeModel {
@@ -198,6 +200,15 @@ func (this *Codes2) GetETFCodes(limit ...int) []string {
 	return this.GetETFs(limit...).Codes()
 }
 
+func (this *Codes2) GetIndexes(limit ...int) CodeModels {
+	size := conv.Default(this.etfs.Len(), limit...)
+	return CodeModels(this.indexes.Limit(size))
+}
+
+func (this *Codes2) GetIndexCodes(limit ...int) []string {
+	return this.GetIndexes(limit...).Codes()
+}
+
 func (this *Codes2) updated() (bool, error) {
 	update := new(UpdateModel)
 	{ //查询或者插入一条数据
@@ -240,6 +251,7 @@ func (this *Codes2) Update() error {
 
 	stocks := []*CodeModel(nil)
 	etfs := []*CodeModel(nil)
+	indexes := []*CodeModel(nil)
 	for _, v := range codes {
 		fullCode := v.FullCode()
 		this.m.Set(fullCode, v)
@@ -248,11 +260,14 @@ func (this *Codes2) Update() error {
 			stocks = append(stocks, v)
 		case protocol.IsETF(fullCode):
 			etfs = append(etfs, v)
+		case protocol.IsIndex(fullCode):
+			indexes = append(indexes, v)
 		}
 	}
 
 	this.stocks = stocks
 	this.etfs = etfs
+	this.indexes = indexes
 	this.all = codes
 
 	return nil

protocol/unit.go

@@ -3,13 +3,14 @@ package protocol
 import (
 	"bytes"
 	"fmt"
-	"github.com/injoyai/conv"
-	"golang.org/x/text/encoding/simplifiedchinese"
-	"golang.org/x/text/transform"
 	"io"
 	"math"
 	"strings"
 	"time"
+
+	"github.com/injoyai/conv"
+	"golang.org/x/text/encoding/simplifiedchinese"
+	"golang.org/x/text/transform"
 )
 
 // String 字节先转小端,再转字符
@@ -291,6 +292,21 @@ func IsETF(code string) bool {
 	return false
 }
 
+// IsIndex 是否是指数,sh000001,sz399001,bj899100
+func IsIndex(code string) bool {
+	if len(code) != 8 {
+		return false
+	}
+	code = strings.ToLower(code)
+	switch {
+	case code[0:2] == ExchangeSH.String() && code[2:5] == "000":
+		return true
+	case code[0:2] == ExchangeSZ.String() && code[2:5] == "399":
+	case code[0:2] == ExchangeBJ.String() && code[2:5] == "899":
+	}
+	return false
+}
+
 // AddPrefix 添加股票/基金代码前缀,针对股票/基金生效,例如000001,会增加前缀sz000001(平安银行),而不是sh000001(上证指数)
 func AddPrefix(code string) string {
 	if len(code) == 6 {