⚡️ 优化 rule.FirstValueInList

plugin_nativesetu/data.go

@@ -5,6 +5,7 @@ import (
 	"image"
 	"io/fs"
 	"os"
+	"strings"
 	"sync"
 
 	"github.com/corona10/goimagehash"
@@ -24,11 +25,7 @@ type setuclass struct {
 	Path  string `db:"path"`  // Path 图片路径
 }
 
-var (
-	setuclasses []string
-	db          = &sql.Sqlite{DBPath: dbfile}
-	mu          sync.RWMutex
-)
+var ns *nsetu
 
 func init() {
 	go func() {
@@ -44,24 +41,33 @@ func init() {
 				logrus.Println("[nsetu] set setu dir to", setupath)
 			}
 		}
-		if file.IsExist(dbfile) {
-			err := db.Open()
-			if err == nil {
-				setuclasses, err = db.ListTables()
-			}
-			if err != nil {
-				logrus.Errorln("[nsetu]", err)
-			}
-		}
+		ns = &nsetu{db: &sql.Sqlite{DBPath: dbfile}}
 	}()
 }
 
-func scanall(path string) error {
-	setuclasses = nil
+type nsetu struct {
+	db *sql.Sqlite
+	mu sync.RWMutex
+}
+
+func (n *nsetu) List() (l []string) {
+	if file.IsExist(n.db.DBPath) {
+		err := n.db.Open()
+		if err == nil {
+			l, err = n.db.ListTables()
+		}
+		if err != nil {
+			logrus.Errorln("[nsetu]", err)
+		}
+	}
+	return
+}
+
+func (n *nsetu) scanall(path string) error {
 	model := &setuclass{}
 	root := os.DirFS(path)
-	_ = db.Close()
-	_ = os.Remove(dbfile)
+	_ = n.db.Close()
+	_ = os.Remove(n.db.DBPath)
 	return fs.WalkDir(root, ".", func(path string, d fs.DirEntry, err error) error {
 		if err != nil {
 			return err
@@ -69,12 +75,11 @@ func scanall(path string) error {
 		if d.IsDir() {
 			clsn := d.Name()
 			if clsn != "." {
-				mu.Lock()
-				err = db.Create(clsn, model)
-				setuclasses = append(setuclasses, clsn)
-				mu.Unlock()
+				n.mu.Lock()
+				err = n.db.Create(clsn, model)
+				n.mu.Unlock()
 				if err == nil {
-					err = scanclass(root, path, clsn)
+					err = n.scanclass(root, path, clsn)
 					if err != nil {
 						logrus.Errorln("[nsetu]", err)
 						return err
@@ -86,17 +91,21 @@ func scanall(path string) error {
 	})
 }
 
-func scanclass(root fs.FS, path, clsn string) error {
+func (n *nsetu) scanclass(root fs.FS, path, clsn string) error {
 	ds, err := fs.ReadDir(root, path)
 	if err != nil {
 		return err
 	}
-	mu.Lock()
-	_ = db.Truncate(clsn)
-	mu.Unlock()
+	n.mu.Lock()
+	_ = n.db.Truncate(clsn)
+	n.mu.Unlock()
 	for _, d := range ds {
-		if !d.IsDir() {
-			relpath := path + "/" + d.Name()
+		nm := d.Name()
+		ln := strings.ToLower(nm)
+		if !d.IsDir() &&
+			(strings.HasSuffix(ln, ".jpg") || strings.HasSuffix(ln, ".jpeg") ||
+				strings.HasSuffix(ln, ".png") || strings.HasSuffix(ln, ".gif") || strings.HasSuffix(ln, ".webp")) {
+			relpath := path + "/" + nm
 			logrus.Debugln("[nsetu] read", relpath)
 			f, e := fs.ReadFile(root, relpath)
 			if e != nil {
@@ -112,10 +121,10 @@ func scanclass(root fs.FS, path, clsn string) error {
 				return e
 			}
 			dhi := int64(dh.GetHash())
-			logrus.Debugln("[nsetu] insert", d.Name(), "with id", dhi, "into", clsn)
-			mu.Lock()
-			err = db.Insert(clsn, &setuclass{ImgID: dhi, Name: d.Name(), Path: relpath})
-			mu.Unlock()
+			logrus.Debugln("[nsetu] insert", nm, "with id", dhi, "into", clsn)
+			n.mu.Lock()
+			err = n.db.Insert(clsn, &setuclass{ImgID: dhi, Name: nm, Path: relpath})
+			n.mu.Unlock()
 			if err != nil {
 				return err
 			}

plugin_nativesetu/main.go

@@ -36,13 +36,13 @@ func init() {
 			"- 刷新所有本地setu\n" +
 			"- 所有本地setu分类",
 	})
-	engine.OnRegex(`^本地(.*)$`, func(ctx *zero.Ctx) bool { return rule.FirstValueInList(setuclasses)(ctx) }).SetBlock(true).
+	engine.OnRegex(`^本地(.*)$`, rule.FirstValueInList(ns)).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			imgtype := ctx.State["regex_matched"].([]string)[1]
 			sc := new(setuclass)
-			mu.RLock()
-			err := db.Pick(imgtype, sc)
-			mu.RUnlock()
+			ns.mu.RLock()
+			err := ns.db.Pick(imgtype, sc)
+			ns.mu.RUnlock()
 			if err != nil {
 				ctx.SendChain(message.Text("ERROR: ", err))
 			} else {
@@ -50,10 +50,10 @@ func init() {
 				ctx.SendChain(message.Text(imgtype, ": ", sc.Name, "\n"), message.Image(p))
 			}
 		})
-	engine.OnRegex(`^刷新本地(.*)$`, func(ctx *zero.Ctx) bool { return rule.FirstValueInList(setuclasses)(ctx) }, zero.SuperUserPermission).SetBlock(true).
+	engine.OnRegex(`^刷新本地(.*)$`, rule.FirstValueInList(ns), zero.SuperUserPermission).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			imgtype := ctx.State["regex_matched"].([]string)[1]
-			err := scanclass(os.DirFS(setupath), imgtype, imgtype)
+			err := ns.scanclass(os.DirFS(setupath), imgtype, imgtype)
 			if err == nil {
 				ctx.SendChain(message.Text("成功！"))
 			} else {
@@ -72,7 +72,7 @@ func init() {
 		})
 	engine.OnFullMatch("刷新所有本地setu", zero.SuperUserPermission).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
-			err := scanall(setupath)
+			err := ns.scanall(setupath)
 			if err == nil {
 				ctx.SendChain(message.Text("成功！"))
 			} else {
@@ -82,9 +82,9 @@ func init() {
 	engine.OnFullMatch("所有本地setu分类").SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			msg := "所有本地setu分类"
-			mu.RLock()
-			for i, c := range setuclasses {
-				n, err := db.Count(c)
+			ns.mu.RLock()
+			for i, c := range ns.List() {
+				n, err := ns.db.Count(c)
 				if err == nil {
 					msg += fmt.Sprintf("\n%02d. %s(%d)", i, c, n)
 				} else {
@@ -92,7 +92,7 @@ func init() {
 					logrus.Errorln("[nsetu]", err)
 				}
 			}
-			mu.RUnlock()
+			ns.mu.RUnlock()
 			ctx.SendChain(message.Text(msg))
 		})
 }