✨ 支持每日推送用户名分遗产情况

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 "还没有订阅任何用户名"

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()]

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]:

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)

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

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)