🔥 ✨ 优化 setutime

Author: fumiama

plugin_setutime/setu_geter.go

@@ -2,6 +2,7 @@
 package setutime
 
 import (
+	"errors"
 	"fmt"
 	"strconv"
 	"strings"
@@ -26,43 +27,24 @@ import (
 
 // Pools 图片缓冲池
 type imgpool struct {
-	Lock  sync.Mutex
-	DB    *sql.Sqlite
-	Path  string
-	Group int64
-	List  []string
-	Max   int
-	Pool  map[string][]*pixiv.Illust
-	Form  int64
+	lock sync.Mutex
+	db   *sql.Sqlite
+	path string
+	max  int
+	pool map[string][]*pixiv.Illust
 }
 
-// NewPoolsCache 返回一个缓冲池对象
-func newPools() *imgpool {
-	cache := &imgpool{
-		DB:    &sql.Sqlite{DBPath: "data/SetuTime/SetuTime.db"},
-		Path:  pixiv.CacheDir,
-		Group: 0,
-		List:  []string{"涩图", "二次元", "风景", "车万"}, // 可以自己加类别，得自己加图片进数据库
-		Max:   10,
-		Pool:  map[string][]*pixiv.Illust{},
-		Form:  0,
-	}
-	// 如果数据库不存在则下载
-	_, _ = fileutil.GetLazyData(cache.DB.DBPath, false, false)
-	for i := range cache.List {
-		if err := cache.DB.Create(cache.List[i], &pixiv.Illust{}); err != nil {
-			panic(err)
-		}
+func (i *imgpool) list() (l []string) {
+	var err error
+	l, err = i.db.ListTables()
+	if err != nil {
+		l = []string{"涩图", "二次元", "风景", "车万"}
 	}
-	return cache
+	return l
 }
 
-var (
-	pool  *imgpool
-	limit = rate.NewManager(time.Minute*1, 5)
-)
-
 func init() { // 插件主体
+	limit := rate.NewManager(time.Minute*1, 5)
 	engine := control.Register("setutime", order.PrioSetuTime, &control.Options{
 		DisableOnDefault: false,
 		Help: "涩图\n" +
@@ -72,29 +54,35 @@ func init() { // 插件主体
 			"- >setu status",
 	})
 	process.SleepAbout1sTo2s()
-	pool = newPools()
-	engine.OnRegex(`^来份(.*)$`, rule.FirstValueInList(pool.List)).SetBlock(true).
+	pool := func() *imgpool {
+		cache := &imgpool{
+			db:   &sql.Sqlite{DBPath: "data/SetuTime/SetuTime.db"},
+			path: pixiv.CacheDir,
+			max:  10,
+			pool: map[string][]*pixiv.Illust{},
+		}
+		// 如果数据库不存在则下载
+		_, _ = fileutil.GetLazyData(cache.db.DBPath, false, false)
+		err := cache.db.Open()
+		if err != nil {
+			panic(err)
+		}
+		for _, imgtype := range cache.list() {
+			if err := cache.db.Create(imgtype, &pixiv.Illust{}); err != nil {
+				panic(err)
+			}
+		}
+		return cache
+	}()
+	engine.OnRegex(`^来份(.*)$`, rule.FirstValueInList(pool.list())).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			if !limit.Load(ctx.Event.UserID).Acquire() {
 				ctx.SendChain(message.Text("请稍后重试0x0..."))
 				return
 			}
 			var imgtype = ctx.State["regex_matched"].([]string)[1]
 			// 补充池子
-			go func() {
-				times := math.Min(pool.Max-pool.size(imgtype), 2)
-				for i := 0; i < times; i++ {
-					illust := &pixiv.Illust{}
-					// 查询出一张图片
-					if err := pool.DB.Pick(imgtype, illust); err != nil {
-						ctx.SendChain(message.Text("ERROR: ", err))
-						continue
-					}
-					// 向缓冲池添加一张图片
-					pool.push(ctx, imgtype, illust)
-					process.SleepAbout1sTo2s()
-				}
-			}()
+			go pool.fill(ctx, imgtype)
 			// 如果没有缓存，阻塞10秒
 			if pool.size(imgtype) == 0 {
 				ctx.SendChain(message.Text("INFO: 正在填充弹药......"))
@@ -105,50 +93,33 @@ func init() { // 插件主体
 				}
 			}
 			// 从缓冲池里抽一张
-			if id := ctx.SendChain(message.Image(file(pool.pop(imgtype)))); id.ID() == 0 {
+			if id := ctx.SendChain(message.Image(pool.popfile(imgtype))); id.ID() == 0 {
 				ctx.SendChain(message.Text("ERROR: 可能被风控了"))
 			}
 		})
 
-	engine.OnRegex(`^添加(.*?)(\d+)$`, rule.FirstValueInList(pool.List), zero.SuperUserPermission).SetBlock(true).
+	engine.OnRegex(`^添加(.*?)(\d+)$`, zero.SuperUserPermission).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			var (
 				imgtype = ctx.State["regex_matched"].([]string)[1]
 				id, _   = strconv.ParseInt(ctx.State["regex_matched"].([]string)[2], 10, 64)
 			)
-			ctx.SendChain(message.Text("少女祈祷中......"))
-			// 查询P站插图信息
-			illust, err := pixiv.Works(id)
+			err := pool.add(ctx, imgtype, id)
 			if err != nil {
-				ctx.SendChain(message.Text("ERROR: ", err))
-				return
-			}
-			// 下载插画
-			if _, err := illust.DownloadToCache(0, strconv.FormatInt(id, 10)+"_p0"); err != nil {
-				ctx.SendChain(message.Text("ERROR: ", err))
-				return
-			}
-			// 发送到发送者
-			if id := ctx.SendChain(message.Image(file(illust))); id.ID() == 0 {
-				ctx.SendChain(message.Text("ERROR: 可能被风控，发送失败"))
+				ctx.SendChain(message.Text("ERROR:", err))
 				return
 			}
-			// 添加插画到对应的数据库table
-			if err := pool.DB.Insert(imgtype, illust); err != nil {
-				ctx.SendChain(message.Text("ERROR: ", err))
-				return
-			}
-			ctx.SendChain(message.Text("添加成功"))
+			ctx.SendChain(message.Text("成功向分类", imgtype, "添加图片", id))
 		})
 
-	engine.OnRegex(`^删除(.*?)(\d+)$`, rule.FirstValueInList(pool.List), zero.SuperUserPermission).SetBlock(true).
+	engine.OnRegex(`^删除(.*?)(\d+)$`, rule.FirstValueInList(pool.list()), zero.SuperUserPermission).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			var (
 				imgtype = ctx.State["regex_matched"].([]string)[1]
 				id, _   = strconv.ParseInt(ctx.State["regex_matched"].([]string)[2], 10, 64)
 			)
 			// 查询数据库
-			if err := pool.DB.Del(imgtype, fmt.Sprintf("WHERE pid=%d", id)); err != nil {
+			if err := pool.remove(imgtype, id); err != nil {
 				ctx.SendChain(message.Text("ERROR: ", err))
 				return
 			}
@@ -159,13 +130,13 @@ func init() { // 插件主体
 	engine.OnFullMatchGroup([]string{">setu status"}).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			state := []string{"[SetuTime]"}
-			for i := range pool.List {
-				num, err := pool.DB.Count(pool.List[i])
+			for _, imgtype := range pool.list() {
+				num, err := pool.db.Count(imgtype)
 				if err != nil {
 					num = 0
 				}
 				state = append(state, "\n")
-				state = append(state, pool.List[i])
+				state = append(state, imgtype)
 				state = append(state, ": ")
 				state = append(state, fmt.Sprintf("%d", num))
 			}
@@ -175,16 +146,9 @@ func init() { // 插件主体
 
 // size 返回缓冲池指定类型的现有大小
 func (p *imgpool) size(imgtype string) int {
-	return len(p.Pool[imgtype])
+	return len(p.pool[imgtype])
 }
 
-/*
-// isFull 返回缓冲池指定类型是否已满
-func (p *imgpool) isFull(imgtype string) bool {
-	return len(p.Pool[imgtype]) >= p.Max
-}*/
-
-// push 向缓冲池插入一张图片
 func (p *imgpool) push(ctx *zero.Ctx, imgtype string, illust *pixiv.Illust) {
 	u := illust.ImageUrls[0]
 	n := u[strings.LastIndex(u, "/")+1 : len(u)-4]
@@ -199,31 +163,30 @@ func (p *imgpool) push(ctx *zero.Ctx, imgtype string, illust *pixiv.Illust) {
 		m.SetFile(fileutil.BOTPATH + "/" + f)
 		_, _ = m.Push(ctxext.SendToSelf(ctx), ctxext.GetMessage(ctx))
 	}
-	p.Lock.Lock()
-	p.Pool[imgtype] = append(p.Pool[imgtype], illust)
-	p.Lock.Unlock()
+	p.lock.Lock()
+	p.pool[imgtype] = append(p.pool[imgtype], illust)
+	p.lock.Unlock()
 }
 
-// Push 在缓冲池拿出一张图片
 func (p *imgpool) pop(imgtype string) (illust *pixiv.Illust) {
-	p.Lock.Lock()
-	defer p.Lock.Unlock()
+	p.lock.Lock()
+	defer p.lock.Unlock()
 	if p.size(imgtype) == 0 {
 		return
 	}
-	illust = p.Pool[imgtype][0]
-	p.Pool[imgtype] = p.Pool[imgtype][1:]
+	illust = p.pool[imgtype][0]
+	p.pool[imgtype] = p.pool[imgtype][1:]
 	return
 }
 
-func file(i *pixiv.Illust) string {
+func (p *imgpool) file(i *pixiv.Illust) string {
 	u := i.ImageUrls[0]
 	m, err := imagepool.GetImage(u[strings.LastIndex(u, "/")+1 : len(u)-4])
 	if err == nil {
 		return m.String()
 	}
 	filename := fmt.Sprint(i.Pid) + "_p0"
-	filepath := fileutil.BOTPATH + `/` + pool.Path + filename
+	filepath := fileutil.BOTPATH + `/` + p.path + filename
 	if fileutil.IsExist(filepath + ".jpg") {
 		return `file:///` + filepath + ".jpg"
 	}
@@ -235,3 +198,52 @@ func file(i *pixiv.Illust) string {
 	}
 	return ""
 }
+
+func (p *imgpool) popfile(imgtype string) string {
+	return p.file(p.pop(imgtype))
+}
+
+// fill 补充池子
+func (p *imgpool) fill(ctx *zero.Ctx, imgtype string) {
+	times := math.Min(p.max-p.size(imgtype), 2)
+	for i := 0; i < times; i++ {
+		illust := &pixiv.Illust{}
+		// 查询出一张图片
+		if err := p.db.Pick(imgtype, illust); err != nil {
+			ctx.SendChain(message.Text("ERROR: ", err))
+			continue
+		}
+		// 向缓冲池添加一张图片
+		p.push(ctx, imgtype, illust)
+		process.SleepAbout1sTo2s()
+	}
+}
+
+func (p *imgpool) add(ctx *zero.Ctx, imgtype string, id int64) error {
+	if err := p.db.Create(imgtype, &pixiv.Illust{}); err != nil {
+		return err
+	}
+	ctx.SendChain(message.Text("少女祈祷中......"))
+	// 查询P站插图信息
+	illust, err := pixiv.Works(id)
+	if err != nil {
+		return err
+	}
+	// 下载插画
+	if _, err := illust.DownloadToCache(0, strconv.FormatInt(id, 10)+"_p0"); err != nil {
+		return err
+	}
+	// 发送到发送者
+	if id := ctx.SendChain(message.Image(p.file(illust))); id.ID() == 0 {
+		return errors.New("可能被风控，发送失败")
+	}
+	// 添加插画到对应的数据库table
+	if err := p.db.Insert(imgtype, illust); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (p *imgpool) remove(imgtype string, id int64) error {
+	return p.db.Del(imgtype, fmt.Sprintf("WHERE pid=%d", id))
+}