细节优化

client.go

@@ -3,6 +3,10 @@ package tdx
 import (
 	"errors"
 	"fmt"
+	"runtime/debug"
+	"sync/atomic"
+	"time"
+
 	"github.com/injoyai/base/maps"
 	"github.com/injoyai/base/maps/wait"
 	"github.com/injoyai/conv"
@@ -10,11 +14,8 @@ import (
 	"github.com/injoyai/ios/client"
 	"github.com/injoyai/ios/module/common"
 	"github.com/injoyai/logs"
-	"github.com/injoyai/tdx/internal/bse"
+	"github.com/injoyai/tdx/lib/bse"
 	"github.com/injoyai/tdx/protocol"
-	"runtime/debug"
-	"sync/atomic"
-	"time"
 )
 
 const (
@@ -560,6 +561,26 @@ func (this *Client) GetIndexAll(Type uint8, code string) (*protocol.KlineResp, e
 	return this.GetIndexUntil(Type, code, func(k *protocol.Kline) bool { return false })
 }
 
+func (this *Client) GetIndexMinute(code string, start, count uint16) (*protocol.KlineResp, error) {
+	return this.GetIndex(protocol.TypeKlineMinute, code, start, count)
+}
+
+func (this *Client) GetIndex5Minute(code string, start, count uint16) (*protocol.KlineResp, error) {
+	return this.GetIndex(protocol.TypeKline5Minute, code, start, count)
+}
+
+func (this *Client) GetIndex15Minute(code string, start, count uint16) (*protocol.KlineResp, error) {
+	return this.GetIndex(protocol.TypeKline15Minute, code, start, count)
+}
+
+func (this *Client) GetIndex30Minute(code string, start, count uint16) (*protocol.KlineResp, error) {
+	return this.GetIndex(protocol.TypeKline30Minute, code, start, count)
+}
+
+func (this *Client) GetIndex60Minute(code string, start, count uint16) (*protocol.KlineResp, error) {
+	return this.GetIndex(protocol.TypeKline60Minute, code, start, count)
+}
+
 func (this *Client) GetIndexDay(code string, start, count uint16) (*protocol.KlineResp, error) {
 	return this.GetIndex(protocol.TypeKlineDay, code, start, count)
 }

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,69 +2,70 @@ package tdx
 
 import (
 	"errors"
+	"iter"
+	"os"
+	"path/filepath"
+	"time"
+
 	"github.com/injoyai/base/maps"
 	"github.com/injoyai/base/types"
 	"github.com/injoyai/conv"
 	"github.com/injoyai/ios"
 	"github.com/injoyai/ios/client"
 	"github.com/injoyai/logs"
-	"github.com/injoyai/tdx/internal/gbbq"
+	"github.com/injoyai/tdx/lib/gbbq"
-	"github.com/injoyai/tdx/internal/xorms"
+	"github.com/injoyai/tdx/lib/xorms"
 	"github.com/injoyai/tdx/protocol"
 	"github.com/robfig/cron/v3"
-	"iter"
-	"os"
-	"path/filepath"
-	"time"
 	"xorm.io/xorm"
 )
 
 type Codes2Option func(*Codes2)
 
-func WithDBFilename(filename string) Codes2Option {
+func WithCodes2Database(filename string) Codes2Option {
 	return func(c *Codes2) {
 		c.dbFilename = filename
 	}
 }
 
-func WithTempDir(dir string) Codes2Option {
+func WithCodes2TempDir(dir string) Codes2Option {
 	return func(c *Codes2) {
 		c.tempDir = dir
 	}
 }
 
-func WithSpec(spec string) Codes2Option {
+func WithCodes2Spec(spec string) Codes2Option {
 	return func(c *Codes2) {
 		c.spec = spec
 	}
 }
 
-func WithUpdateKey(key string) Codes2Option {
+func WithCodes2UpdateKey(key string) Codes2Option {
 	return func(c *Codes2) {
 		c.updateKey = key
 	}
 }
 
-func WithRetry(retry int) Codes2Option {
+func WithCodes2Retry(retry int) Codes2Option {
 	return func(c *Codes2) {
 		c.retry = retry
 	}
 }
 
-func WithClient(c *Client) Codes2Option {
+func WithCodes2Client(c *Client) Codes2Option {
 	return func(cs *Codes2) {
 		cs.c = c
 	}
 }
 
-func WithDial(dial ios.DialFunc, op ...client.Option) Codes2Option {
+func WithCodes2Dial(dial ios.DialFunc, op ...client.Option) Codes2Option {
 	return func(c *Codes2) {
 		c.dial = dial
 		c.dialOption = op
 	}
 }
 
-func WithDialOption(op ...client.Option) Codes2Option {
+func WithCodes2DialOption(op ...client.Option) Codes2Option {
 	return func(c *Codes2) {
 		c.dialOption = op
 	}
@@ -76,7 +77,7 @@ func NewCodes2(op ...Codes2Option) (*Codes2, error) {
 		tempDir:    filepath.Join(DefaultDataDir, "temp"),
 		spec:       "10 0 9 * * *",
 		updateKey:  "codes",
-		retry:      3,
+		retry:      DefaultRetry,
 		dial:       NewRangeDial(Hosts),
 		dialOption: nil,
 		m:          maps.NewGeneric[string, *CodeModel](),
@@ -116,7 +117,7 @@ func NewCodes2(op ...Codes2Option) (*Codes2, error) {
 	// 定时更新
 	cr := cron.New(cron.WithSeconds())
 	_, err = cr.AddFunc(cs.spec, func() {
-		for i := 0; i < 3; i++ {
+		for i := 0; i == 0 || i < cs.retry; i++ {
 			if err := cs.Update(); err != nil {
 				logs.Err(err)
 				<-time.After(time.Minute * 5)
@@ -149,12 +150,13 @@ type Codes2 struct {
 		内部字段
 	*/
 
-	c      *Client                           //
+	c       *Client                           //
-	db     *xorms.Engine                     //
+	db      *xorms.Engine                     //
-	stocks types.List[*CodeModel]            //缓存
+	stocks  types.List[*CodeModel]            //股票缓存
-	etfs   types.List[*CodeModel]            //缓存
+	etfs    types.List[*CodeModel]            //etf缓存
-	all    types.List[*CodeModel]            //缓存
+	indexes types.List[*CodeModel]            //指数缓存
-	m      *maps.Generic[string, *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

example/CheckGetHistoryErr/main.go

@@ -1,13 +1,14 @@
 package main
 
 import (
+	"time"
+
 	"github.com/injoyai/logs"
 	"github.com/injoyai/tdx"
-	"time"
 )
 
 func main() {
-	m, err := tdx.NewManage(nil)
+	m, err := tdx.NewManage()
 	logs.PanicErr(err)
 
 	codes := m.Codes.GetStocks().Codes()

example/GetETFCodeNumber/main.go

@@ -0,0 +1,29 @@
+package main
+
+import (
+	"strings"
+
+	"github.com/injoyai/logs"
+	"github.com/injoyai/tdx"
+)
+
+func main() {
+	cs, err := tdx.NewCodes2()
+	logs.PanicErr(err)
+
+	ls := cs.GetETFCodes()
+
+	shNumber := 0
+	szNumber := 0
+	for _, v := range ls {
+		switch {
+		case strings.HasPrefix(v, "sh"):
+			shNumber++
+		case strings.HasPrefix(v, "sz"):
+			szNumber++
+		}
+	}
+
+	logs.Debug("sh:", shNumber)
+	logs.Debug("sz:", szNumber)
+}

example/Manage/main.go

@@ -0,0 +1,23 @@
+package main
+
+import (
+	"github.com/injoyai/logs"
+	"github.com/injoyai/tdx"
+)
+
+func main() {
+	m, err := tdx.NewManage()
+	logs.PanicErr(err)
+
+	err = m.Do(func(c *tdx.Client) error {
+		resp, err := c.GetIndexDayAll("sh000001")
+		if err != nil {
+			return err
+		}
+		for _, v := range resp.List {
+			logs.Debug(v)
+		}
+		return nil
+	})
+	logs.PanicErr(err)
+}

example/PullKline/main.go

@@ -2,16 +2,17 @@ package main
 
 import (
 	"context"
+	"path/filepath"
+	"time"
+
 	"github.com/injoyai/logs"
 	"github.com/injoyai/tdx"
 	"github.com/injoyai/tdx/extend"
-	"path/filepath"
-	"time"
 )
 
 func main() {
 
-	m, err := tdx.NewManage(nil)
+	m, err := tdx.NewManage()
 	logs.PanicErr(err)
 
 	err = extend.NewPullKline(extend.PullKlineConfig{

example/PullTrade/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+
 	"github.com/injoyai/logs"
 	"github.com/injoyai/tdx"
 	"github.com/injoyai/tdx/extend"
@@ -11,7 +12,7 @@ func main() {
 
 	pt := extend.NewPullTrade("./data/trade")
 
-	m, err := tdx.NewManage(nil)
+	m, err := tdx.NewManage()
 	logs.PanicErr(err)
 
 	err = pt.PullYear(context.Background(), m, 2025, "sz000001")

extend/codes-bj.go

@@ -1,7 +1,7 @@
 package extend
 
 import (
-	"github.com/injoyai/tdx/internal/bse"
+	"github.com/injoyai/tdx/lib/bse"
 )
 
 func GetBjCodes() ([]string, error) {

lib/gbbq/gbbq.go

@@ -12,7 +12,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"github.com/injoyai/base/types"
-	"github.com/injoyai/tdx/internal/zip"
+	"github.com/injoyai/tdx/lib/zip"
 	"io"
 	"math"
 	"net/http"

manage.go

@@ -2,170 +2,224 @@ package tdx
 
 import (
 	"errors"
+	"sync"
+
+	"github.com/injoyai/conv"
 	"github.com/injoyai/ios/client"
 	"github.com/robfig/cron/v3"
-	"sync"
-	"time"
 )
 
 const (
+	DefaultClients     = 1
+	DefaultRetry       = 3
 	DefaultDataDir     = "./data"
 	DefaultDatabaseDir = "./data/database"
 )
 
-func NewManageMysql(cfg *ManageConfig, op ...client.Option) (*Manage, error) {
+func NewManageMysql(op ...Option) (*Manage, error) {
-	//初始化配置
+	return NewManage(
-	if cfg == nil {
+		WithOptions(op...),
-		cfg = &ManageConfig{}
+		WithDialCodes(func(c *Client, database string) (ICodes, error) {
-	}
+			if database == "" {
-	if cfg.CodesFilename == "" {
+				return nil, errors.New("未配置Codes的数据库")
-		return nil, errors.New("未配置Codes的数据库")
+			}
-	}
+			return NewCodesMysql(c, database)
-	if cfg.WorkdayFileName == "" {
+		}),
-		return nil, errors.New("未配置Workday的数据库")
+		WithDialWorkday(func(c *Client, database string) (*Workday, error) {
-	}
+			if database == "" {
-	if cfg.Dial == nil {
+				return nil, errors.New("未配置Workday的数据库")
-		cfg.Dial = DialDefault
+			}
-	}
+			return NewWorkdayMysql(c, database)
-
+		}),
-	//通用客户端
+	)
-	commonClient, err := cfg.Dial(op...)
-	if err != nil {
-		return nil, err
-	}
-	commonClient.Wait.SetTimeout(time.Second * 5)
-
-	//代码管理
-	codes, err := NewCodesMysql(commonClient, cfg.CodesFilename)
-	if err != nil {
-		return nil, err
-	}
-
-	//工作日管理
-	workday, err := NewWorkdayMysql(commonClient, cfg.WorkdayFileName)
-	if err != nil {
-		return nil, err
-	}
-
-	//连接池
-	p, err := NewPool(func() (*Client, error) {
-		return cfg.Dial(op...)
-	}, cfg.Number)
-	if err != nil {
-		return nil, err
-	}
-
-	return &Manage{
-		Pool:    p,
-		Config:  cfg,
-		Codes:   codes,
-		Workday: workday,
-	}, nil
 }
 
-func NewManage(cfg *ManageConfig, op ...client.Option) (*Manage, error) {
+func NewManageSqlite(op ...Option) (*Manage, error) {
-	//初始化配置
+	return NewManage(
-	if cfg == nil {
+		WithCodesDatabase(DefaultDatabaseDir+"/codes.db"),
-		cfg = &ManageConfig{}
+		WithWorkdayDatabase(DefaultDatabaseDir+"/workday.db"),
-	}
+		WithOptions(op...),
-	if cfg.CodesFilename == "" {
+		WithDialCodes(func(c *Client, database string) (ICodes, error) {
-		cfg.CodesFilename = DefaultDatabaseDir + "/codes.db"
+			return NewCodesSqlite(c, database)
-	}
+		}),
-	if cfg.WorkdayFileName == "" {
+		WithDialWorkday(func(c *Client, database string) (*Workday, error) {
-		cfg.WorkdayFileName = DefaultDatabaseDir + "/workday.db"
+			return NewWorkdaySqlite(c, database)
-	}
+		}),
-	if cfg.Dial == nil {
+	)
-		cfg.Dial = DialDefault
-	}
-
-	//通用客户端
-	commonClient, err := cfg.Dial(op...)
-	if err != nil {
-		return nil, err
-	}
-	commonClient.Wait.SetTimeout(time.Second * 5)
-
-	//代码管理
-	codes, err := NewCodesSqlite(commonClient, cfg.CodesFilename)
-	if err != nil {
-		return nil, err
-	}
-
-	//工作日管理
-	workday, err := NewWorkdaySqlite(commonClient, cfg.WorkdayFileName)
-	if err != nil {
-		return nil, err
-	}
-
-	//连接池
-	p, err := NewPool(func() (*Client, error) {
-		return cfg.Dial(op...)
-	}, cfg.Number)
-	if err != nil {
-		return nil, err
-	}
-
-	return &Manage{
-		Pool:    p,
-		Config:  cfg,
-		Codes:   codes,
-		Workday: workday,
-	}, nil
 }
 
-func NewManage2(cfg *ManageConfig, op ...client.Option) (*Manage, error) {
+func NewManageSqlite2(op ...Option) (*Manage, error) {
-	//初始化配置
+	return NewManage(
-	if cfg == nil {
+		WithCodesDatabase(DefaultDatabaseDir+"/codes2.db"),
-		cfg = &ManageConfig{}
+		WithWorkdayDatabase(DefaultDatabaseDir+"/workday.db"),
-	}
+		WithOptions(op...),
-	if cfg.CodesFilename == "" {
+		WithDialCodes(func(c *Client, database string) (ICodes, error) {
-		cfg.CodesFilename = DefaultDatabaseDir + "/codes2.db"
+			return NewCodes2(
-	}
+				WithCodes2Client(c),
-	if cfg.WorkdayFileName == "" {
+				WithCodes2Database(database),
-		cfg.WorkdayFileName = DefaultDatabaseDir + "/workday.db"
+			)
-	}
+		}),
-	if cfg.Dial == nil {
+		WithDialWorkday(func(c *Client, database string) (*Workday, error) {
-		cfg.Dial = DialDefault
+			return NewWorkdaySqlite(c, database)
+		}),
+	)
+
+}
+
+func NewManage(op ...Option) (m *Manage, err error) {
+
+	m = &Manage{
+		clients:         DefaultClients,
+		dial:            DialDefault,
+		dialOptions:     nil,
+		dialCodes:       nil,
+		codesDatabase:   DefaultDatabaseDir + "/codes2.db",
+		dialWorkday:     nil,
+		workdayDatabase: DefaultDatabaseDir + "/workday.db",
+		Pool:            nil,
+		Codes:           nil,
+		Workday:         nil,
+		cron:            nil,
+		once:            sync.Once{},
 	}
 
-	//通用客户端
+	for _, v := range op {
-	commonClient, err := cfg.Dial(op...)
+		if v != nil {
+			v(m)
+		}
+	}
+
+	m.clients = conv.Select(m.clients <= 0, 1, m.clients)
+	m.dial = conv.Select(m.dial == nil, DialDefault, m.dial)
+
+	//连接池
+	m.Pool, err = NewPool(func() (*Client, error) { return m.dial(m.dialOptions...) }, m.clients)
 	if err != nil {
 		return nil, err
 	}
-	commonClient.Wait.SetTimeout(time.Second * 5)
 
 	//代码管理
-	codes, err := NewCodes2(WithClient(commonClient), WithDBFilename(cfg.CodesFilename))
+	if m.Codes == nil {
-	if err != nil {
+		if m.dialCodes == nil {
-		return nil, err
+			m.dialCodes = func(c *Client, database string) (ICodes, error) {
+				return NewCodes2(WithCodes2Client(c), WithCodes2Database(database))
+			}
+		}
+		err = m.Pool.Do(func(c *Client) error {
+			m.Codes, err = m.dialCodes(c, m.codesDatabase)
+			return err
+		})
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	//工作日管理
-	workday, err := NewWorkdaySqlite(commonClient, cfg.WorkdayFileName)
+	if m.Workday == nil {
-	if err != nil {
+		if m.dialWorkday == nil {
-		return nil, err
+			m.dialWorkday = func(c *Client, database string) (*Workday, error) {
+				return NewWorkdaySqlite(c, database)
+			}
+		}
+		err = m.Pool.Do(func(c *Client) error {
+			m.Workday, err = m.dialWorkday(c, m.workdayDatabase)
+			return err
+		})
+		if err != nil {
+			return nil, err
+		}
 	}
 
-	//连接池
+	return
-	p, err := NewPool(func() (*Client, error) {
+}
-		return cfg.Dial(op...)
-	}, cfg.Number)
-	if err != nil {
-		return nil, err
-	}
 
-	return &Manage{
+/*
-		Pool:    p,
+
-		Config:  cfg,
+
-		Codes:   codes,
+
-		Workday: workday,
+ */
-	}, nil
+
+type Option func(m *Manage)
+type DialWorkdayFunc func(c *Client, database string) (*Workday, error)
+type DialCodesFunc func(c *Client, database string) (ICodes, error)
+
+func WithClients(clients int) Option {
+	return func(m *Manage) {
+		m.clients = clients
+	}
+}
+
+func WithDial(dial func(op ...client.Option) (*Client, error), op ...client.Option) Option {
+	return func(m *Manage) {
+		m.dial = dial
+		m.dialOptions = op
+	}
+}
+
+func WithDialOptions(op ...client.Option) Option {
+	return func(m *Manage) {
+		m.dialOptions = op
+	}
+}
+
+func WithCodes(codes ICodes) Option {
+	return func(m *Manage) {
+		m.Codes = codes
+	}
+}
+
+func WithDialCodes(dial DialCodesFunc) Option {
+	return func(m *Manage) {
+		m.dialCodes = dial
+	}
+}
+
+func WithCodesDatabase(database string) Option {
+	return func(m *Manage) {
+		m.codesDatabase = database
+	}
+}
+
+func WithWorkday(w *Workday) Option {
+	return func(m *Manage) {
+		m.Workday = w
+	}
+}
+
+func WithDialWorkday(dial DialWorkdayFunc) Option {
+	return func(m *Manage) {
+		m.dialWorkday = dial
+	}
+}
+
+func WithWorkdayDatabase(database string) Option {
+	return func(m *Manage) {
+		m.workdayDatabase = database
+	}
+}
+
+func WithOptions(op ...Option) Option {
+	return func(m *Manage) {
+		for _, v := range op {
+			v(m)
+		}
+	}
 }
 
 type Manage struct {
+	clients         int
+	dial            func(op ...client.Option) (cli *Client, err error)
+	dialOptions     []client.Option
+	dialCodes       func(c *Client, database string) (ICodes, error)
+	codesDatabase   string
+	dialWorkday     DialWorkdayFunc
+	workdayDatabase string
+
+	/*
+
+	 */
+
 	*Pool
-	Config  *ManageConfig
 	Codes   ICodes
 	Workday *Workday
 	cron    *cron.Cron
@@ -186,6 +240,13 @@ func (this *Manage) RangeETFs(f func(code string)) {
 	}
 }
 
+// RangeIndexes 遍历所有指数
+func (this *Manage) RangeIndexes(f func(code string)) {
+	for _, v := range this.Codes.GetETFs() {
+		f(v.FullCode())
+	}
+}
+
 // AddWorkdayTask 添加工作日任务
 func (this *Manage) AddWorkdayTask(spec string, f func(m *Manage)) {
 	this.once.Do(func() {

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 字节先转小端,再转字符
@@ -285,12 +286,27 @@ func IsETF(code string) bool {
 		return true
 
 	case code[0:2] == ExchangeSZ.String() &&
-		(code[2:4] == "15" || code[2:4] == "16"):
+		(code[2:4] == "15"):
 		return true
 	}
 	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 {

workday.go

@@ -2,6 +2,11 @@ package tdx
 
 import (
 	"errors"
+	"iter"
+	"os"
+	"path/filepath"
+	"time"
+
 	_ "github.com/glebarez/go-sqlite"
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/injoyai/base/maps"
@@ -10,9 +15,6 @@ import (
 	"github.com/injoyai/logs"
 	"github.com/injoyai/tdx/protocol"
 	"github.com/robfig/cron/v3"
-	"os"
-	"path/filepath"
-	"time"
 	"xorm.io/core"
 	"xorm.io/xorm"
 )
@@ -154,7 +156,7 @@ func (this *Workday) TodayIs() bool {
 func (this *Workday) RangeYear(year int, f func(t time.Time) bool) {
 	this.Range(
 		time.Date(year, 1, 1, 0, 0, 0, 0, time.Local),
-		time.Date(year, 12, 31, 0, 0, 0, 0, time.Local),
+		time.Date(year, 12, 31, 0, 0, 0, 1, time.Local),
 		f,
 	)
 }
@@ -162,8 +164,6 @@ func (this *Workday) RangeYear(year int, f func(t time.Time) bool) {
 // Range 遍历指定范围的工作日,推荐start带上时间15:00,这样当天小于15点不会触发
 func (this *Workday) Range(start, end time.Time, f func(t time.Time) bool) {
 	start = conv.Select(start.Before(protocol.ExchangeEstablish), protocol.ExchangeEstablish, start)
-	//now := IntegerDay(time.Now())
-	//end = conv.Select(end.After(now), now, end).Add(1)
 	for ; start.Before(end); start = start.Add(time.Hour * 24) {
 		if this.Is(start) {
 			if !f(start) {
@@ -173,13 +173,36 @@ func (this *Workday) Range(start, end time.Time, f func(t time.Time) bool) {
 	}
 }
 
-// RangeDesc 倒序遍历工作日,从今天-1990年12月19日(上海交易所成立时间)
+func (this *Workday) IterYear(year int, desc ...bool) iter.Seq[time.Time] {
-func (this *Workday) RangeDesc(f func(t time.Time) bool) {
+	return this.Iter(
-	t := IntegerDay(time.Now())
+		time.Date(year, 1, 1, 0, 0, 0, 0, time.Local),
-	for ; t.After(time.Date(1990, 12, 18, 0, 0, 0, 0, time.Local)); t = t.Add(-time.Hour * 24) {
+		time.Date(year, 12, 31, 0, 0, 0, 1, time.Local),
-		if this.Is(t) {
+		desc...,
-			if !f(t) {
+	)
-				return
+}
+
+// Iter 遍历指定范围的工作日,推荐start带上时间15:00,这样当天小于15点不会触发
+func (this *Workday) Iter(start, end time.Time, desc ...bool) iter.Seq[time.Time] {
+	start = conv.Select(start.Before(protocol.ExchangeEstablish), protocol.ExchangeEstablish, start)
+	if len(desc) > 0 && desc[0] {
+		//倒序遍历
+		return func(yield func(time.Time) bool) {
+			for ; end.After(start); end = end.Add(-time.Hour * 24) {
+				if this.Is(end) {
+					if !yield(end) {
+						return
+					}
+				}
+			}
+		}
+	}
+	//正序遍历
+	return func(yield func(time.Time) bool) {
+		for ; start.Before(end); start = start.Add(time.Hour * 24) {
+			if this.Is(start) {
+				if !yield(start) {
+					return
+				}
 			}
 		}
 	}