wyj/starbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: Some model modifications and code optimization

Browse Source
This commit is contained in:
LWR 2022-11-23 01:07:26 +08:00
parent 35bc6a4a21
commit dc6125f142
5 changed files with 66 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
starbot/core/bot.py
View File

@ -81,20 +81,8 @@ class StarBot:
await redis.hset("LiveStatus", up.room_id, status)
await redis.hset("StartTime", up.room_id, base["live_time"])
await redis.hset_ifnotexists("EndTime", up.room_id, 0)
await redis.hset_ifnotexists("RoomDanmuCount", up.room_id, 0)
await redis.hset_ifnotexists("RoomDanmuTotal", up.room_id, 0)
await redis.hset_ifnotexists("RoomBoxCount", up.room_id, 0)
await redis.hset_ifnotexists("RoomBoxTotal", up.room_id, 0)
await redis.hset_ifnotexists("RoomBoxProfit", up.room_id, 0)
await redis.hset_ifnotexists("RoomBoxProfitTotal", up.room_id, 0)
await redis.hset_ifnotexists("RoomGiftProfit", up.room_id, 0)
await redis.hset_ifnotexists("RoomGiftTotal", up.room_id, 0)
await redis.hset_ifnotexists("RoomScProfit", up.room_id, 0)
await redis.hset_ifnotexists("RoomScTotal", up.room_id, 0)
await redis.hset_ifnotexists("RoomGuardCount", up.room_id, "0-0-0")
await redis.hset_ifnotexists("RoomGuardTotal", up.room_id, "0-0-0")
# 连接直播间
for up in self.__datasource.get_up_list():
try:
await up.connect()
@ -105,7 +93,7 @@ class StarBot:
if config.get("USE_HTTP_API"):
asyncio.get_event_loop().create_task(http_init(self.__datasource))
# 启动 Bot
# 启动消息推送模块
if not self.__datasource.bots:
logger.error("不存在需要启动的 Bot 账号, 请先在数据源中配置完毕后再重新运行")
return
@ -113,8 +101,6 @@ class StarBot:
Ariadne.options["default_account"] = self.__datasource.bots[0].qq
logger.info("开始运行 Ariadne 消息推送模块")
logger.disable("graia.ariadne.service")
logger.disable("launart")
for bot in self.__datasource.bots:
bot.start_sender()
@ -141,6 +127,9 @@ class StarBot:
)
logger.remove()
logger.add(sys.stderr, format=logger_format, level="INFO")
logger.disable("graia.ariadne.model")
logger.disable("graia.ariadne.service")
logger.disable("launart")
bcc = create(Broadcast)
loop = bcc.loop

9
starbot/core/datasource.py
View File

@ -42,7 +42,7 @@ class DataSource(metaclass=abc.ABCMeta):
Raises:
DataSourceException: 配置中包含重复 uid
"""
self.__up_list = [x for up in map(lambda bot: bot.ups, self.bots) for x in up]
self.__up_list = [x for up in map(lambda b: b.ups, self.bots) for x in up]
self.__up_map = dict(zip(map(lambda up: up.uid, self.__up_list), self.__up_list))
self.__uid_list = list(self.__up_map.keys())
if len(set(self.__uid_list)) < len(self.__uid_list):
@ -195,7 +195,6 @@ class MySQLDataSource(DataSource):
self.__port = port or int(config.get("MYSQL_PORT"))
self.__db = db or config.get("MYSQL_DB")
self.__pool: Optional[aiomysql.pool.Pool] = None
self.__loop = asyncio.get_event_loop()
async def __connect(self):
"""
@ -205,6 +204,7 @@ class MySQLDataSource(DataSource):
DataSourceException: 连接数据库失败
"""
try:
self.__loop = asyncio.get_event_loop()
self.__pool = await aiomysql.create_pool(host=self.__host,
port=self.__port,
user=self.__username,
@ -265,7 +265,7 @@ class MySQLDataSource(DataSource):
uid = now_user.get("uid")
live_on = await self.__query(
"SELECT g.`uid`, g.`uname`, g.`room_id`, `key`, `type`, `num`, `enabled`, `at_all`, `message` "
"SELECT g.`uid`, g.`uname`, g.`room_id`, `key`, `type`, `num`, `enabled`, `message` "
"FROM `groups` AS `g` LEFT JOIN `live_on` AS `l` "
"ON g.`uid` = l.`uid` AND g.`index` = l.`index` "
f"WHERE g.`uid` = {uid} "
@ -295,9 +295,8 @@ class MySQLDataSource(DataSource):
targets = []
for i, target in enumerate(live_on):
if all((live_on[i]["enabled"], live_on[i]["at_all"], live_on[i]["message"])):
if all((live_on[i]["enabled"], live_on[i]["message"])):
on = LiveOn(enabled=live_on[i]["enabled"],
at_all=live_on[i]["at_all"],
message=live_on[i]["message"])
else:
on = LiveOn()

8
starbot/core/easter_egg.py
View File

@ -1,8 +0,0 @@
s = "Unccl oveguqnl gb zr!"
d = {}
for c in (65, 97):
for i in range(26):
d[chr(i+c)] = chr((i+13) % 26 + c)
print("".join([d.get(c, c) for c in s]))

36
starbot/core/model.py
View File

@ -6,6 +6,8 @@ from graia.ariadne.message.chain import MessageChain
from graia.ariadne.message.element import Plain, At, AtAll, Image
from pydantic import BaseModel, PrivateAttr
from ..exception import DataSourceException
class LiveOn(BaseModel):
"""
@ -20,13 +22,11 @@ class LiveOn(BaseModel):
enabled: Optional[bool] = False
"""是否启用开播推送。默认：False"""
at_all: Optional[bool] = False
"""是否 @全体成员。默认：False"""
message: Optional[str] = ""
"""
开播推送内容模板。
可用占位符：{uname}主播昵称，{title}直播间标题，{url}直播间链接，{cover}直播间封面图，{next}消息分条。
专用占位符：{uname} 主播昵称，{title} 直播间标题，{url} 直播间链接，{cover} 直播间封面图。
通用占位符：{next} 消息分条，{atall} @全体成员，{at114514} @指定QQ号，{urlpic=链接} 网络图片，{pathpic=路径} 本地图片。
默认：""
"""
@ -34,12 +34,12 @@ class LiveOn(BaseModel):
def default(cls):
"""
获取功能全部开启的默认 LiveOn 实例
默认配置：启用开播推送，启用 @全体成员，推送内容模板为 "{uname} 正在直播 {title}\n{url}{next}{cover}"
默认配置：启用开播推送，推送内容模板为 "{uname} 正在直播 {title}\n{url}{next}{cover}"
"""
return LiveOn(enabled=True, at_all=True, message=LiveOn.DEFAULT_MESSAGE)
return LiveOn(enabled=True, message=LiveOn.DEFAULT_MESSAGE)
def __str__(self):
return f"启用: {self.enabled}\n@全体: {self.at_all}\n推送内容:\n{self.message}"
return f"启用: {self.enabled}\n推送内容:\n{self.message}"
class LiveOff(BaseModel):
@ -49,7 +49,7 @@ class LiveOff(BaseModel):
或使用 LiveOff.default() 获取功能全部开启的默认配置
"""
DEFAULT_MESSAGE: Optional[str] = "{uname} 直播结束了\n{time}{next}{danmu_count}{danmu_mvp}{box_profit}"
DEFAULT_MESSAGE: Optional[str] = "{uname} 直播结束了"
"""默认消息模板"""
enabled: Optional[bool] = False
@ -58,7 +58,8 @@ class LiveOff(BaseModel):
message: Optional[str] = ""
"""
下播推送内容模板。
可用占位符：{uname}主播昵称，{time}本次直播时长，{danmu_count}弹幕总数，{danmu_mvp}弹幕MVP，{box_profit}宝盒盈亏，{next}消息分条。
专用占位符：{uname}主播昵称。
通用占位符：{next} 消息分条，{atall} @全体成员，{at114514} @指定QQ号，{urlpic=链接} 网络图片，{pathpic=路径} 本地图片。
默认：""
"""
@ -115,7 +116,8 @@ class DynamicUpdate(BaseModel):
message: Optional[str] = ""
"""
动态推送内容模板。
可用占位符：{uname}主播昵称，{action}动态操作类型（发表了新动态，转发了新动态，投稿了新视频...），{url}动态链接（若为发表视频、专栏等则为视频、专栏等对应的链接）。
专用占位符：{uname}主播昵称，{action}动态操作类型（发表了新动态，转发了新动态，投稿了新视频...），{url}动态链接（若为发表视频、专栏等则为视频、专栏等对应的链接）。
通用占位符：{next} 消息分条，{atall} @全体成员，{at114514} @指定QQ号，{urlpic=链接} 网络图片，{pathpic=路径} 本地图片。
默认：""
"""
@ -174,6 +176,15 @@ class PushTarget(BaseModel):
super().__init__(**data)
if not self.key:
self.key = "-".join([str(self.id), str(self.type.value)])
self.__raise_for_not_invalid_placeholders()
def __raise_for_not_invalid_placeholders(self):
"""
使用不合法的占位符时抛出异常
"""
if self.type == PushType.Friend:
if "{at" in self.live_on.message or "{at" in self.live_off.message or "{at" in self.dynamic_update.message:
raise DataSourceException(f"好友类型的推送目标 (QQ: {self.id}) 推送内容中不能含有 @ 消息, 请检查配置后重试")
def __eq__(self, other):
if isinstance(other, PushTarget):
@ -267,6 +278,7 @@ class Message(BaseModel):
chain.append(Plain(msg[:code_end + 1]))
msg = msg[code_end + 1:]
next_code = msg.find("{")
chains.append(chain)
if len(chain) != 0:
chains.append(chain)
return chains
return chains

46
starbot/utils/redis.py
View File

@ -1,4 +1,4 @@
from typing import Any, Union
from typing import Any, Union, Optional, List
import aioredis
from loguru import logger
@ -27,12 +27,14 @@ async def init():
# String
async def get(key: str) -> str:
return str(await __redis.get(key))
async def delete(key: str):
await __redis.delete(key)
async def geti(key: str) -> int:
return int(await __redis.get(key))
# List
async def rpush(key: str, value: Any):
await __redis.rpush(key, value)
# Hash
@ -41,18 +43,36 @@ async def hexists(key: str, hkey: Union[str, int]) -> bool:
return await __redis.hexists(key, hkey)
async def hget(key: str, hkey: Union[str, int]) -> str:
return str(await __redis.hget(key, hkey))
async def hgeti(key: str, hkey: Union[str, int]) -> int:
return int(await __redis.hget(key, hkey))
result = await __redis.hget(key, hkey)
if result is None:
return 0
return int(result)
async def hgetf1(key: str, hkey: Union[str, int]) -> float:
result = await __redis.hget(key, hkey)
if result is None:
return 0.0
return float("{:.1f}".format(float(result)))
async def hset(key: str, hkey: Union[str, int], value: Any):
await __redis.hset(key, hkey, value)
async def hset_ifnotexists(key: str, hkey: Union[str, int], value: Any):
if not await hexists(key, hkey):
await hset(key, hkey, value)
async def hincrby(key: str, hkey: Union[str, int], value: Optional[int] = 1) -> int:
return await __redis.hincrby(key, hkey, value)
async def hincrbyfloat(key: str, hkey: Union[str, int], value: Optional[float] = 1.0) -> float:
return await __redis.hincrbyfloat(key, hkey, value)
# Zset
async def zunionstore(dest: str, source: Union[str, List[str]]):
if isinstance(source, str):
await __redis.zunionstore(dest, [dest, source])
if isinstance(source, list):
await __redis.zunionstore(dest, source)