53 lines
1.3 KiB
Python
53 lines
1.3 KiB
Python
from io import BytesIO
|
|
|
|
from pyrogram.types import Story
|
|
|
|
from userbot import PyroClient
|
|
from userbot.modules.stories import repository
|
|
|
|
|
|
def _peer_id(story: Story) -> int:
|
|
if story.chat is not None:
|
|
return story.chat.id or 0
|
|
if story.from_user is not None:
|
|
return story.from_user.id or 0
|
|
return 0
|
|
|
|
|
|
@PyroClient.on_story()
|
|
async def on_story(client: PyroClient, story: Story) -> None:
|
|
ctx = client.capture
|
|
if ctx is None:
|
|
return
|
|
media_kind = story.media.name.lower() if story.media else None
|
|
storage_key: str | None = None
|
|
file_size: int | None = None
|
|
downloaded = False
|
|
if not story.deleted and story.media is not None:
|
|
buffer = await client.download_media(story, in_memory=True)
|
|
if isinstance(buffer, BytesIO):
|
|
data = buffer.getvalue()
|
|
storage_key = ctx.storage.put(data)
|
|
file_size = len(data)
|
|
downloaded = True
|
|
await repository.upsert_story(
|
|
ctx.pool,
|
|
ctx.account_id,
|
|
_peer_id(story),
|
|
story.id,
|
|
story.date,
|
|
story.expire_date,
|
|
story.caption,
|
|
media_kind,
|
|
storage_key,
|
|
file_size,
|
|
story.views,
|
|
str(story.raw),
|
|
pinned=bool(story.pinned),
|
|
deleted=bool(story.deleted),
|
|
downloaded=downloaded,
|
|
)
|
|
|
|
|
|
handlers = on_story.handlers
|