57 lines
1.3 KiB
Python
57 lines
1.3 KiB
Python
import asyncio
|
|
import contextlib
|
|
from io import BytesIO
|
|
|
|
from pyrogram import Client
|
|
from pyrogram.types import (
|
|
Animation,
|
|
Audio,
|
|
Document,
|
|
Message,
|
|
Photo,
|
|
Sticker,
|
|
Story,
|
|
Video,
|
|
VideoNote,
|
|
Voice,
|
|
)
|
|
|
|
from utils.logging import logger
|
|
|
|
type Downloadable = (
|
|
str
|
|
| Message
|
|
| Story
|
|
| Audio
|
|
| Document
|
|
| Photo
|
|
| Sticker
|
|
| Animation
|
|
| Video
|
|
| Voice
|
|
| VideoNote
|
|
)
|
|
|
|
STOP_TIMEOUT = 5.0
|
|
|
|
_reset_lock = asyncio.Lock()
|
|
|
|
|
|
async def _reset_media_sessions(client: Client) -> None:
|
|
async with _reset_lock:
|
|
sessions = list(client.media_sessions.values())
|
|
client.media_sessions.clear()
|
|
for session in sessions:
|
|
with contextlib.suppress(Exception):
|
|
await asyncio.wait_for(session.stop(), STOP_TIMEOUT)
|
|
logger.warning(f"[yellow]Dropped {len(sessions)} stale media session(s).[/]")
|
|
|
|
|
|
async def download_bytes(client: Client, target: Downloadable) -> bytes | None:
|
|
try:
|
|
buffer = await client.download_media(target, in_memory=True)
|
|
except (TimeoutError, OSError):
|
|
await _reset_media_sessions(client)
|
|
buffer = await client.download_media(target, in_memory=True)
|
|
return buffer.getvalue() if isinstance(buffer, BytesIO) else None
|