优化codes-server,待完成

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 injoyai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

example/Codes2/main.go

@@ -2,6 +2,8 @@ package main
 
 import (
 	"fmt"
+	"strings"
+
 	"github.com/injoyai/logs"
 	"github.com/injoyai/tdx"
 )
@@ -13,4 +15,11 @@ func main() {
 	c := cs.Get("sz000001")
 
 	fmt.Println(c.FloatStock, c.TotalStock)
+
+	for _, v := range cs.GetIndexes().Codes() {
+		if strings.HasPrefix(v, "sz") {
+			logs.Debug(v)
+		}
+	}
+
 }

extend/codes-server.go

@@ -3,43 +3,147 @@ package extend
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/injoyai/conv"
-	"github.com/injoyai/tdx"
 	"io"
+	"iter"
 	"net/http"
-	"path/filepath"
+
+	"github.com/injoyai/base/maps"
+	"github.com/injoyai/conv"
+	"github.com/injoyai/logs"
+	"github.com/injoyai/tdx"
+	"github.com/robfig/cron/v3"
 )
 
-func ListenCodesHTTP(port int, filename ...string) error {
-	code, err := tdx.DialCodes(conv.Default(filepath.Join(tdx.DefaultDatabaseDir, "codes.db"), filename...))
+func ListenCodesHTTP(port int, op ...tdx.Codes2Option) error {
+	code, err := tdx.NewCodes2(op...)
 	if err != nil {
 		return nil
 	}
+	succ := func(w http.ResponseWriter, data any) {
+		w.WriteHeader(http.StatusOK)
+		w.Write(conv.Bytes(data))
+	}
 	return http.ListenAndServe(fmt.Sprintf(":%d", port), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.RequestURI {
+		case "/all":
+
 		case "/stocks":
-			ls := code.GetStocks()
-			w.WriteHeader(http.StatusOK)
-			w.Write(conv.Bytes(ls))
+			succ(w, code.GetStocks())
 		case "/etfs":
-			ls := code.GetETFs()
-			w.WriteHeader(http.StatusOK)
-			w.Write(conv.Bytes(ls))
+			succ(w, code.GetETFs())
+		case "/indexes":
+			succ(w, code.GetIndexes())
 		default:
 			http.NotFound(w, r)
 		}
 	}))
 }
 
-func DialCodesHTTP(address string) *CodesHTTP {
-	return &CodesHTTP{address: address}
+func DialCodesHTTP(address string) (c *CodesHTTP, err error) {
+	c = &CodesHTTP{address: address}
+	cr := cron.New(cron.WithSeconds())
+	_, err = cr.AddFunc("0 20 9 * * *", func() { logs.PrintErr(c.Update()) })
+	if err != nil {
+		return
+	}
+	err = c.Update()
+	if err != nil {
+		return
+	}
+	cr.Start()
+	return c, nil
 }
 
 type CodesHTTP struct {
 	address string
+	stocks  tdx.CodeModels
+	etfs    tdx.CodeModels
+	indexes tdx.CodeModels
+	m       maps.Generic[string, *tdx.CodeModel]
 }
 
-func (this *CodesHTTP) getList(path string) ([]string, error) {
+func (this *CodesHTTP) Iter() iter.Seq2[string, *tdx.CodeModel] {
+	return func(yield func(string, *tdx.CodeModel) bool) {
+		for _, v := range this.stocks {
+			if !yield(v.FullCode(), v) {
+				return
+			}
+		}
+		for _, v := range this.etfs {
+			if !yield(v.FullCode(), v) {
+				return
+			}
+		}
+		for _, v := range this.indexes {
+			if !yield(v.FullCode(), v) {
+				return
+			}
+		}
+	}
+}
+
+func (this *CodesHTTP) Get(code string) *tdx.CodeModel {
+	return this.m.MustGet(code)
+}
+
+func (this *CodesHTTP) GetName(code string) string {
+	v := this.m.MustGet(code)
+	if v != nil {
+		return v.Name
+	}
+	return ""
+}
+
+func (this *CodesHTTP) GetStocks(limit ...int) tdx.CodeModels {
+	return this.stocks
+}
+
+func (this *CodesHTTP) GetStockCodes(limit ...int) []string {
+	return this.stocks.Codes()
+}
+
+func (this *CodesHTTP) GetETFs(limit ...int) tdx.CodeModels {
+	return this.etfs
+}
+
+func (this *CodesHTTP) GetETFCodes(limit ...int) []string {
+	return this.etfs.Codes()
+}
+
+func (this *CodesHTTP) GetIndexes(limits ...int) tdx.CodeModels {
+	return this.indexes
+}
+
+func (this *CodesHTTP) GetIndexCodes(limits ...int) []string {
+	return this.indexes.Codes()
+}
+
+func (this *CodesHTTP) Update() (err error) {
+	this.stocks, err = this.getList("/stocks")
+	if err != nil {
+		return
+	}
+	for _, v := range this.stocks {
+		this.m.Set(v.FullCode(), v)
+	}
+	this.etfs, err = this.getList("/etfs")
+	if err != nil {
+		return
+	}
+	for _, v := range this.etfs {
+		this.m.Set(v.FullCode(), v)
+	}
+	this.indexes, err = this.getList("/indexes")
+	if err != nil {
+		return
+	}
+	for _, v := range this.indexes {
+		this.m.Set(v.FullCode(), v)
+	}
+	return
+}
+
+func (this *CodesHTTP) getList(path string) (tdx.CodeModels, error) {
 	resp, err := http.DefaultClient.Get(this.address + path)
 	if err != nil {
 		return nil, err
@@ -52,15 +156,7 @@ func (this *CodesHTTP) getList(path string) ([]string, error) {
 	if err != nil {
 		return nil, err
 	}
-	ls := []string(nil)
+	ls := tdx.CodeModels{}
 	err = json.Unmarshal(bs, &ls)
 	return ls, err
 }
-
-func (this *CodesHTTP) GetStocks() ([]string, error) {
-	return this.getList("/stocks")
-}
-
-func (this *CodesHTTP) GetETFs() ([]string, error) {
-	return this.getList("/etfs")
-}

protocol/unit.go

@@ -302,7 +302,9 @@ func IsIndex(code string) bool {
 	case code[0:2] == ExchangeSH.String() && code[2:5] == "000":
 		return true
 	case code[0:2] == ExchangeSZ.String() && code[2:5] == "399":
+		return true
 	case code[0:2] == ExchangeBJ.String() && code[2:5] == "899":
+		return true
 	}
 	return false
 }