260 lines
8.5 KiB
Python
260 lines
8.5 KiB
Python
"""Progress tracking and throttled Telegram progress messages.
|
||
|
||
Transfer callbacks only update in-memory counters (cheap, thread-safe enough for
|
||
plain attribute writes). A background task renders the board every
|
||
PROGRESS_INTERVAL seconds and edits one Telegram message per chat, so progress never
|
||
floods the chat with new messages.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import html
|
||
import logging
|
||
import time
|
||
from collections import deque
|
||
from dataclasses import dataclass, field
|
||
from typing import Callable, Protocol
|
||
|
||
log = logging.getLogger(__name__)
|
||
|
||
|
||
class HtmlReply(str):
|
||
"""A message whose text is Telegram HTML (dynamic parts already escaped)."""
|
||
|
||
|
||
def esc(value) -> str:
|
||
return html.escape(str(value), quote=False)
|
||
|
||
|
||
class Notifier(Protocol):
|
||
async def send(self, chat_id: int, text: str) -> int | None: ...
|
||
|
||
async def edit(self, chat_id: int, message_id: int, text: str) -> None: ...
|
||
|
||
|
||
class LogNotifier:
|
||
"""Notifier used when no control bot is configured (headless mode)."""
|
||
|
||
async def send(self, chat_id: int, text: str) -> int | None:
|
||
log.info("%s", text.replace("\n", " | "))
|
||
return None
|
||
|
||
async def edit(self, chat_id: int, message_id: int, text: str) -> None:
|
||
log.info("%s", text.replace("\n", " | "))
|
||
|
||
|
||
def format_bytes(value: float | None) -> str:
|
||
if value is None:
|
||
return "?"
|
||
for unit in ("B", "KB", "MB", "GB"):
|
||
if abs(value) < 1024:
|
||
return f"{value:.0f} {unit}" if unit == "B" else f"{value:.1f} {unit}"
|
||
value /= 1024
|
||
return f"{value:.2f} TB"
|
||
|
||
|
||
def format_duration(seconds: float | None) -> str:
|
||
if seconds is None:
|
||
return "?"
|
||
seconds = int(round(seconds))
|
||
hours, rem = divmod(seconds, 3600)
|
||
minutes, secs = divmod(rem, 60)
|
||
if hours:
|
||
return f"{hours}h {minutes}m"
|
||
if minutes:
|
||
return f"{minutes}m {secs}s"
|
||
return f"{secs}s"
|
||
|
||
|
||
@dataclass
|
||
class TransferProgress:
|
||
total: int | None
|
||
done: int = 0
|
||
clock: Callable[[], float] = time.monotonic
|
||
window: float = 30.0
|
||
_samples: deque = field(default_factory=deque)
|
||
|
||
def update(self, done: int, total: int | None = None) -> None:
|
||
if total is not None:
|
||
self.total = total
|
||
self.done = done
|
||
now = self.clock()
|
||
self._samples.append((now, done))
|
||
while len(self._samples) > 2 and now - self._samples[0][0] > self.window:
|
||
self._samples.popleft()
|
||
|
||
@property
|
||
def percent(self) -> float | None:
|
||
if not self.total:
|
||
return None
|
||
return min(100.0, 100.0 * self.done / self.total)
|
||
|
||
@property
|
||
def speed(self) -> float | None:
|
||
if len(self._samples) < 2:
|
||
return None
|
||
(t0, b0), (t1, b1) = self._samples[0], self._samples[-1]
|
||
if t1 - t0 < 1.0 or b1 < b0:
|
||
return None
|
||
return (b1 - b0) / (t1 - t0)
|
||
|
||
@property
|
||
def eta(self) -> float | None:
|
||
speed = self.speed
|
||
if not speed or not self.total:
|
||
return None
|
||
return max(0.0, (self.total - self.done) / speed)
|
||
|
||
def describe(self) -> str:
|
||
pct = self.percent
|
||
parts = [f"{pct:.0f}%" if pct is not None else format_bytes(self.done)]
|
||
parts.append(f"({format_bytes(self.done)} / {format_bytes(self.total)})")
|
||
text = " ".join(parts)
|
||
if self.speed:
|
||
text += f" · {format_bytes(self.speed)}/s"
|
||
if self.eta is not None and self.done < (self.total or 0):
|
||
text += f" · ETA {format_duration(self.eta)}"
|
||
return text
|
||
|
||
|
||
@dataclass
|
||
class ProgressBoard:
|
||
position: int = 0
|
||
total: int = 0
|
||
topic: str = ""
|
||
title: str = ""
|
||
status: str = "Idle"
|
||
download: TransferProgress | None = None
|
||
upload: TransferProgress | None = None
|
||
clock: Callable[[], float] = time.monotonic
|
||
# Set while an /upload task runs: id, start, end, current, completed, skipped, failed
|
||
task: dict | None = None
|
||
|
||
def begin_item(self, position: int, total: int, topic: str, title: str) -> None:
|
||
self.position, self.total, self.topic, self.title = position, total, topic, title
|
||
self.status = "Preparing"
|
||
self.download = None
|
||
self.upload = None
|
||
|
||
def start_download(self, total: int | None) -> None:
|
||
self.status = "Downloading"
|
||
self.download = TransferProgress(total, clock=self.clock)
|
||
|
||
def start_upload(self, total: int) -> None:
|
||
self.status = "Uploading"
|
||
self.upload = TransferProgress(total, clock=self.clock)
|
||
|
||
_TASK_STATUS = {
|
||
"Downloading": "Downloading from Telegram",
|
||
"Uploading": "Uploading to Google Drive",
|
||
"Verifying": "Verifying upload in Google Drive",
|
||
}
|
||
|
||
def render(self) -> str:
|
||
if self.task is not None:
|
||
return self._render_task()
|
||
if self.status == "Idle":
|
||
return "Queue idle"
|
||
lines = [
|
||
f"Processing video {self.position} of {self.total}",
|
||
f"Topic: {self.topic}",
|
||
f"Title: {self.title}",
|
||
f"Status: {self.status}",
|
||
f"Download: {self.download.describe() if self.download else 'Waiting'}",
|
||
f"Upload: {self.upload.describe() if self.upload else 'Waiting'}",
|
||
]
|
||
return "\n".join(lines)
|
||
|
||
def _render_task(self) -> str:
|
||
t = self.task or {}
|
||
status = self._TASK_STATUS.get(self.status, self.status)
|
||
lines = [
|
||
"🚀 <b>Upload Task Running</b>",
|
||
"",
|
||
f"Task: #{esc(t.get('id'))}",
|
||
f"Range: {esc(t.get('start'))}–{esc(t.get('end'))}",
|
||
f"Current Message: {esc(t.get('current') or '-')}",
|
||
f"Current Video: {esc(self.title or '-')}",
|
||
f"Topic: {esc(self.topic or '-')}",
|
||
f"Status: {esc(status)}",
|
||
]
|
||
if self.download:
|
||
lines.append(f"Download: {esc(self.download.describe())}")
|
||
if self.upload:
|
||
lines.append(f"Upload: {esc(self.upload.describe())}")
|
||
lines += [
|
||
"",
|
||
f"Completed: {t.get('completed', 0)}",
|
||
f"Skipped: {t.get('skipped', 0)}",
|
||
f"Failed: {t.get('failed', 0)}",
|
||
]
|
||
return HtmlReply("\n".join(lines))
|
||
|
||
|
||
class ProgressReporter:
|
||
def __init__(
|
||
self,
|
||
notifier: Notifier,
|
||
chat_ids: tuple[int, ...] | list[int],
|
||
interval: float,
|
||
*,
|
||
clock: Callable[[], float] = time.monotonic,
|
||
):
|
||
self.notifier = notifier
|
||
self.chat_ids = tuple(chat_ids) or (0,)
|
||
self.interval = interval
|
||
self.board = ProgressBoard(clock=clock)
|
||
self._message_ids: dict[int, int] = {}
|
||
self._last_text: str | None = None
|
||
self._task: asyncio.Task | None = None
|
||
|
||
async def start(self) -> None:
|
||
self._message_ids.clear()
|
||
self._last_text = None
|
||
if self._task is None or self._task.done():
|
||
self._task = asyncio.create_task(self._loop(), name="progress-reporter")
|
||
|
||
async def stop(self) -> None:
|
||
if self._task:
|
||
self._task.cancel()
|
||
try:
|
||
await self._task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
self._task = None
|
||
await self.refresh()
|
||
|
||
async def _loop(self) -> None:
|
||
while True:
|
||
await asyncio.sleep(self.interval)
|
||
await self.refresh()
|
||
|
||
async def refresh(self) -> None:
|
||
text = self.board.render()
|
||
if text == self._last_text:
|
||
return
|
||
self._last_text = text
|
||
for chat_id in self.chat_ids:
|
||
try:
|
||
message_id = self._message_ids.get(chat_id)
|
||
if message_id is None:
|
||
sent = await self.notifier.send(chat_id, text)
|
||
if sent is not None:
|
||
self._message_ids[chat_id] = sent
|
||
else:
|
||
await self.notifier.edit(chat_id, message_id, text)
|
||
except Exception as exc: # progress must never break processing
|
||
log.debug("Progress update to %s failed: %s", chat_id, exc)
|
||
|
||
def new_message(self) -> None:
|
||
"""Start a fresh progress message (e.g. for each /upload task)."""
|
||
self._message_ids.clear()
|
||
self._last_text = None
|
||
|
||
async def notify(self, text: str, *, extra_chat_ids: tuple[int, ...] = ()) -> None:
|
||
for chat_id in dict.fromkeys((*self.chat_ids, *extra_chat_ids)):
|
||
try:
|
||
await self.notifier.send(chat_id, text)
|
||
except Exception as exc:
|
||
log.warning("Notification to %s failed: %s", chat_id, exc)
|