From 0b5b1a5f8000d76e97a695345ced3009cb163c1a Mon Sep 17 00:00:00 2001 From: xtaodada Date: Fri, 4 Nov 2022 23:54:19 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20=E6=94=AF=E6=8C=81=E6=AF=8F?= =?UTF-8?q?=E6=97=A5=E6=8E=A8=E9=80=81=E7=94=A8=E6=88=B7=E5=90=8D=E5=88=86?= =?UTF-8?q?=E9=81=97=E4=BA=A7=E6=83=85=E5=86=B5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- defs/fragment.py | 22 +++++++++++++- models/fragment.py | 55 +++++++++++++++++++++++++++++++-- models/lofter.py | 5 +-- models/models/fragment.py | 9 ++++++ models/sqlite.py | 3 +- modules/fragment.py | 64 +++++++++++++++++++++++++++++++++++++-- 6 files changed, 147 insertions(+), 11 deletions(-) create mode 100644 models/models/fragment.py diff --git a/defs/fragment.py b/defs/fragment.py index 2e7399c..360259a 100644 --- a/defs/fragment.py +++ b/defs/fragment.py @@ -1,10 +1,11 @@ import contextlib import re from datetime import datetime +from typing import Optional from bs4 import BeautifulSoup from init import request -from models.fragment import AuctionStatus, UserName, TON_TO_USD_RATE, Price +from models.fragment import AuctionStatus, UserName, TON_TO_USD_RATE, Price, FragmentSubText, FragmentSub class NotAvailable(Exception): @@ -83,3 +84,22 @@ async def parse_fragment(username: str) -> UserName: except AssertionError: html = await search_fragment_html(username) return search_user(username, html) + + +async def parse_sub(status: FragmentSubText, user: Optional[UserName], cid: int) -> str: + if status == FragmentSubText.Subscribe: + if user.status == [AuctionStatus.Sold, AuctionStatus.Unavailable]: + return "用户名已被卖出或者已被注册，无法订阅" + if await FragmentSub.get_by_cid_and_username(cid, user.name): + return "已经订阅过了这个用户名" + await FragmentSub.subscribe(cid, user.name) + return "订阅成功" + elif status == FragmentSubText.Unsubscribe: + if data := (await FragmentSub.get_by_cid_and_username(cid, user.name)): + await FragmentSub.unsubscribe(data) + return "取消订阅成功" + return "当前没有订阅这个用户名" + elif status == FragmentSubText.List: + if data := (await FragmentSub.get_by_cid(cid)): + return "目前已订阅：\n\n" + "\n".join([f"{i+1}. @{d.username}" for i, d in enumerate(data)]) + return "还没有订阅任何用户名" diff --git a/models/fragment.py b/models/fragment.py index 309891a..1823fd7 100644 --- a/models/fragment.py +++ b/models/fragment.py @@ -1,8 +1,12 @@ from datetime import datetime, timezone, timedelta from enum import Enum -from typing import Optional - +from typing import Optional, cast, List from pydantic import BaseModel +from sqlalchemy import select +from sqlmodel.ext.asyncio.session import AsyncSession + +from init import sqlite +from models.models.fragment import Fragment TON_TO_USD_RATE = {"rate": 1.61} @@ -85,3 +89,50 @@ class UserName(BaseModel): text += f"售价：{self.now_price.text}\n" \ f"距离出售结束：{self.end_human_time}\n" return text + + +class FragmentSubText(Enum): + Subscribe = "订阅" + Unsubscribe = "退订" + List = "订阅列表" + + +class FragmentSub: + @staticmethod + async def subscribe(cid: int, username: str): + async with sqlite.Session() as session: + session = cast(AsyncSession, session) + data = Fragment(cid=cid, username=username) + session.add(data) + await session.commit() + + @staticmethod + async def unsubscribe(data: Fragment): + async with sqlite.Session() as session: + session = cast(AsyncSession, session) + await session.delete(data) + await session.commit() + + @staticmethod + async def get_by_cid_and_username(cid: int, username: str) -> Optional[Fragment]: + async with sqlite.Session() as session: + session = cast(AsyncSession, session) + statement = select(Fragment).where(Fragment.cid == cid and Fragment.username == username) + results = await session.exec(statement) + return post[0] if (post := results.first()) else None + + @staticmethod + async def get_by_cid(cid: int) -> List[Fragment]: + async with sqlite.Session() as session: + session = cast(AsyncSession, session) + statement = select(Fragment).where(Fragment.cid == cid) + results = await session.exec(statement) + return [item[0] for item in results.all()] + + @staticmethod + async def get_all() -> List[Fragment]: + async with sqlite.Session() as session: + session = cast(AsyncSession, session) + statement = select(Fragment) + results = await session.exec(statement) + return [item[0] for item in results.all()] diff --git a/models/lofter.py b/models/lofter.py index e2d9ba4..9259f0c 100644 --- a/models/lofter.py +++ b/models/lofter.py @@ -18,10 +18,7 @@ class LofterPost: check = Lofter.post_id == post_id statement = select(Lofter).where(check) results = await session.exec(statement) - if post := results.first(): - return post[0] - else: - return None + return post[0] if (post := results.first()) else None @staticmethod async def get_by_post_id(post_id: str) -> Optional[Lofter]: diff --git a/models/models/fragment.py b/models/models/fragment.py new file mode 100644 index 0000000..311586b --- /dev/null +++ b/models/models/fragment.py @@ -0,0 +1,9 @@ +from sqlmodel import SQLModel, Field + + +class Fragment(SQLModel, table=True): + __table_args__ = dict(mysql_charset='utf8mb4', mysql_collate="utf8mb4_general_ci") + + cid: int = Field(primary_key=True) + username: str = Field(primary_key=True) + timestamp: int = Field(default=0) diff --git a/models/sqlite.py b/models/sqlite.py index 1b920ae..e5aa0aa 100644 --- a/models/sqlite.py +++ b/models/sqlite.py @@ -1,8 +1,9 @@ from sqlmodel import SQLModel from models.models.lofter import Lofter +from models.models.fragment import Fragment -__all__ = ["Lofter"] +__all__ = ["Lofter", "Fragment"] from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.orm import sessionmaker diff --git a/modules/fragment.py b/modules/fragment.py index 58dd972..bd9eced 100644 --- a/modules/fragment.py +++ b/modules/fragment.py @@ -1,14 +1,58 @@ +import contextlib import re -from pyrogram import Client +from pyrogram import Client, filters, ContinuePropagation +from pyrogram.enums import ChatMemberStatus from pyrogram.types import InlineQuery, InlineQueryResultArticle, InputTextMessageContent, InlineKeyboardMarkup, \ - InlineKeyboardButton + InlineKeyboardButton, Message -from defs.fragment import parse_fragment, NotAvailable +from models.fragment import FragmentSubText, FragmentSub, AuctionStatus +from defs.fragment import parse_fragment, NotAvailable, parse_sub +from init import user_me, bot +from scheduler import scheduler, add_delete_message_job QUERY_PATTERN = re.compile(r"^@\w[a-zA-Z0-9_]{3,32}$") +@Client.on_message(filters.incoming & filters.command(["username", f"username@{user_me.username}"])) +async def fragment_command(client: Client, message: Message): + status = None + user = None + if len(message.command) <= 1: + return await message.reply("没有找到要查询的用户名 ...") + elif message.command[1] == "订阅列表": + status = FragmentSubText.List + elif len(message.command) > 2: + if message.command[2] not in ["订阅", "退订"]: + return await message.reply("只能查询一个用户名 ...") + status = FragmentSubText(message.command[2]) + if status and message.from_user: + data = await client.get_chat_member(message.chat.id, message.from_user.id) + if data.status not in [ChatMemberStatus.ADMINISTRATOR, ChatMemberStatus.OWNER]: + rep = await message.reply("You are not an admin of this chat.") + add_delete_message_job(rep) + raise ContinuePropagation + if status == FragmentSubText.List: + text = await parse_sub(status, user, message.chat.id) + else: + username = message.command[1] + if not username.startswith("@"): + username = f"@{username}" + if not QUERY_PATTERN.match(username): + return await message.reply("无效的用户名") + username = username[1:] + try: + user = await parse_fragment(username) + text = user.text + except NotAvailable: + text = "解析失败了 ... 请稍后再试" + except Exception: + text = "查询失败了 ... 请稍后再试" + if status and user is not None: + text = await parse_sub(status, user, message.chat.id) + await message.reply(text) + + @Client.on_inline_query() async def fragment_inline(_, inline_query: InlineQuery): username = inline_query.query @@ -59,3 +103,17 @@ async def fragment_inline(_, inline_query: InlineQuery): switch_pm_parameter="start", cache_time=0 ) + + +@scheduler.scheduled_job("cron", hour="8", minute="1", id="fragment.sub") +async def fragment_sub() -> None: + data = await FragmentSub.get_all() + if not data: + return + for item in data: + with contextlib.suppress(NotAvailable, Exception): + user = await parse_fragment(item.username) + text = user.text + if user.status in [AuctionStatus.Sold, AuctionStatus.Unavailable]: + await FragmentSub.unsubscribe(item) + await bot.send_message(item.cid, text)