71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
from collections.abc import AsyncIterator, Callable
|
|
|
|
from pyrogram import Client
|
|
from pyrogram.errors import FloodPremiumWait, FloodWait, RPCError
|
|
from pyrogram.types import Story
|
|
|
|
from userbot.modules.capture.context import CaptureContext
|
|
from userbot.modules.jobs.context import JobContext
|
|
from userbot.modules.jobs.registry import register
|
|
from userbot.modules.stories.service import save_story
|
|
|
|
SAVE_EVERY = 10
|
|
|
|
StorySource = Callable[[], AsyncIterator[Story]]
|
|
|
|
|
|
def _sources(client: Client, peer_id: int, *, own: bool) -> dict[str, StorySource]:
|
|
sources: dict[str, StorySource] = {
|
|
"active": lambda: client.get_chat_stories(peer_id),
|
|
"pinned": lambda: client.get_pinned_stories(peer_id),
|
|
}
|
|
if own:
|
|
sources["archived"] = lambda: client.get_archived_stories(peer_id)
|
|
return sources
|
|
|
|
|
|
async def _drain(
|
|
ctx: JobContext, capture: CaptureContext, name: str, source: StorySource
|
|
) -> int:
|
|
client = ctx.client
|
|
if client is None:
|
|
return 0
|
|
saved = 0
|
|
async for story in source():
|
|
try:
|
|
await save_story(client, capture, story)
|
|
except (FloodWait, FloodPremiumWait):
|
|
raise
|
|
except RPCError:
|
|
continue
|
|
saved += 1
|
|
if saved % SAVE_EVERY == 0:
|
|
await ctx.report_progress({"saved": saved, "source": name})
|
|
if await ctx.is_canceled():
|
|
break
|
|
return saved
|
|
|
|
|
|
@register("backfill_stories")
|
|
async def backfill_stories(ctx: JobContext) -> None:
|
|
client = ctx.client
|
|
if client is None:
|
|
return
|
|
capture = getattr(client, "capture", None)
|
|
if capture is None:
|
|
return
|
|
peer_id = ctx.job.params["peer_id"]
|
|
own = client.me is not None and client.me.id == peer_id
|
|
saved = 0
|
|
errors: dict[str, str] = {}
|
|
for name, source in _sources(client, peer_id, own=own).items():
|
|
try:
|
|
saved += await _drain(ctx, capture, name, source)
|
|
except (FloodWait, FloodPremiumWait):
|
|
raise
|
|
except RPCError as exc:
|
|
errors[name] = type(exc).__name__
|
|
if await ctx.is_canceled():
|
|
break
|
|
await ctx.report_progress({"saved": saved, "done": True, "errors": errors})
|