Spaces:
Running
Running
Download bot/dbflow.py from Mega-Evolutions/Auto-Rename: direct link, hf CLI and curl.
- Browser
- Download file 17.9 kB
-
https://huggingface.co/spaces/Mega-Evolutions/Auto-Rename/resolve/main/bot/dbflow.py
- Command line
-
hf download hf://spaces/Mega-Evolutions/Auto-Rename/bot/dbflow.py
-
curl -L -o dbflow.py https://huggingface.co/spaces/Mega-Evolutions/Auto-Rename/resolve/main/bot/dbflow.py
17.9 kB
| """Upload a manually renamed episode to the Database Channel, and have the channel post bot publish it. | |
| All of this exists only while DB_CHANNEL and COMMUNICATION_CHAT are both set. | |
| 1. After /rename with a manual format has sent the renamed episode, ``DbFlow.offer`` asks (Yes / No) whether to upload it. | |
| 2. Yes: the message becomes "Processing…" and the episode is copied into the Database Channel on Telegram's side (forwarded | |
| without the "Forwarded from" line): nothing is downloaded or uploaded again. | |
| 3. When the format has a Channel, "File uploaded to the Database Channel" is then written in the communication chat, with the | |
| episode number, that channel id, the file name and the id of the copy. The channel post bot answers there, replying to it: | |
| "message received", then "post published successfully" or "publication failed: <reason>". | |
| 4. Those answers are read here (``on_chat_message``) and shown in the Processing message. No answer at all: the channel post | |
| bot is down. | |
| A format without a Channel is only uploaded: no notice is written and nothing is published. | |
| Nothing about a running upload is stored: the buttons carry what is needed, and a restart in the middle just leaves the | |
| Processing message as it was. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import logging | |
| import re | |
| import secrets | |
| from dataclasses import dataclass, field | |
| from typing import Any, Dict, Optional, Set, Tuple | |
| from telethon import errors | |
| from . import panels | |
| from .config import Config | |
| from .parser import parse_info, split_ext | |
| from .utils import clip, esc | |
| log = logging.getLogger(__name__) | |
| PREFIX = "dbu" # callback data: dbu:y:<chat>:<message>:<channel or 0> (Yes), dbu:n (No) | |
| ACK_WAIT = 60.0 # seconds the channel post bot has to say "message received" before it counts as down | |
| RESULT_WAIT = 900.0 # ... and then to say whether the post was published (it waits at least 10 s for the Share URL button) | |
| NOTICE_HEAD = "File uploaded to the Database Channel" | |
| _FAILED = re.compile(r"(?i)publication\s+failed\s*[:\-–]?[ \t]*(.*)") | |
| _PUBLISHED = re.compile(r"(?i)post\s+published\s+successfully") | |
| _RECEIVED = re.compile(r"(?i)message\s+received") | |
| def episode_of(file_name: str) -> str: | |
| """The episode number of a file name, as written in it (``05``); empty when it has none.""" | |
| return parse_info(split_ext(file_name or "")[0]).episode | |
| def line_of(label: str, text: str) -> str: | |
| """The value of a ``Label: value`` line of a message; empty when there is no such line.""" | |
| m = re.search(rf"(?im)^[ \t]*{label}[ \t]*:[ \t]*(.+?)[ \t]*$", text or "") | |
| return m.group(1).strip() if m else "" | |
| def classify(text: Optional[str]) -> Tuple[str, str]: | |
| """What a message of the channel post bot says: ``("fail", reason)``, ``("ok", "")``, ``("ack", "")`` or ``("other", "")``.""" | |
| text = text or "" | |
| failed = _FAILED.search(text) | |
| if failed: | |
| return "fail", " ".join(failed.group(1).split()) | |
| if _PUBLISHED.search(text): | |
| return "ok", "" | |
| if _RECEIVED.search(text): | |
| return "ack", "" | |
| return "other", "" | |
| def notice_text(*, episode: str, channel: int, file: str, post_id: Optional[int], ref: str) -> str: | |
| """The message for the communication chat. The channel post bot reads the Episode and the Channel from it.""" | |
| lines = [NOTICE_HEAD] | |
| if episode: | |
| lines.append(f"Episode: {episode}") | |
| lines.append(f"Channel: {channel}") | |
| lines.append(f"File: {file}") | |
| if post_id: | |
| lines.append(f"Post: {post_id}") | |
| lines.append(f"Ref: {ref}") | |
| return "\n".join(lines) | |
| def reply_target(event: Any) -> Optional[int]: | |
| """The id of the message an incoming message replies to.""" | |
| message = getattr(event, "message", None) | |
| return getattr(message, "reply_to_msg_id", None) or getattr(event, "reply_to_msg_id", None) or None | |
| def _name(name: str) -> str: | |
| return f"<code>{esc(clip(name, 120))}</code>" if name else "" | |
| def _why(exc: BaseException) -> str: | |
| return esc(clip(f"{type(exc).__name__}: {exc}", 250)) | |
| def offer_text(name: str, episode: str, channel: Optional[int]) -> str: | |
| lines = ["📤 <b>Upload this episode to the Database Channel?</b>", _name(name)] | |
| if channel: | |
| lines.append(f"🎞 Episode: <b>{esc(episode)}</b>" if episode else "⚠️ No episode number found in the file name (the post uses it).") | |
| lines.append(f"📣 Channel: <code>{channel}</code>: the channel post bot posts it there.") | |
| else: | |
| lines.append("This format has no Channel, so the episode is only uploaded: no post is made.") | |
| return "\n".join(line for line in lines if line) | |
| def processing_text(name: str, *steps: str) -> str: | |
| return "\n".join(line for line in ["⏳ <b>Processing…</b>", _name(name), *steps] if line) | |
| def published_text(name: str, reply: str, channel: int) -> str: | |
| where = line_of("channel", reply) or str(channel) | |
| link = line_of("post", reply) | |
| lines = ["✅ <b>Post published successfully</b>", _name(name), f"📣 {esc(clip(where, 80))}"] | |
| if link: | |
| lines.append(esc(clip(link, 200))) | |
| return "\n".join(line for line in lines if line) | |
| def failed_text(name: str, reason: str) -> str: | |
| lines = ["❌ <b>Publication failed</b>", _name(name), esc(clip(reason or "No reason was given.", 600)), "The episode is in the Database Channel."] | |
| return "\n".join(line for line in lines if line) | |
| def down_text(name: str, seconds: float) -> str: | |
| return "\n".join( | |
| [ | |
| "⚠️ <b>The channel post bot is down or not responding</b>", | |
| _name(name), | |
| f"It did not answer in the communication chat within {int(seconds)} seconds. " | |
| "The episode is in the Database Channel, but nothing was published.", | |
| ] | |
| ) | |
| def silent_text(name: str, seconds: float) -> str: | |
| return "\n".join( | |
| [ | |
| "⚠️ <b>No word from the channel post bot</b>", | |
| _name(name), | |
| f"It got the message but did not say whether the post was published in {int(seconds // 60)} minutes. Check the channel.", | |
| ] | |
| ) | |
| class Flow: | |
| """An upload that waits for the channel post bot's answers.""" | |
| ref: str # in the notice and echoed in every answer: how the answers find their way back | |
| inbox: "asyncio.Queue[Tuple[str, str]]" = field(default_factory=asyncio.Queue) # (ack | ok | fail, the message) | |
| notice_id: int = 0 # the notice in the communication chat; an answer is a reply to it | |
| class DbFlow: | |
| def __init__(self, client: Any, cfg: Config, *, ack_wait: float = ACK_WAIT, result_wait: float = RESULT_WAIT): | |
| self.client = client | |
| self.cfg = cfg | |
| self.ack_wait = ack_wait | |
| self.result_wait = result_wait | |
| self.flows: Dict[str, Flow] = {} | |
| self._busy: Set[Tuple[int, int]] = set() # (chat, offer message) of the uploads under way: a second tap does nothing | |
| self._tasks: Set["asyncio.Future[Any]"] = set() | |
| # ------------------------------------------------------------------ the offer | |
| async def offer(self, job: Any, target: int, sent: Any) -> None: | |
| """Ask whether the episode that was just sent (``sent``, in ``target``) goes to the Database Channel too.""" | |
| sent_id = getattr(sent, "id", None) | |
| if not isinstance(sent_id, int): | |
| return # nothing to point at | |
| channel = (job.fmt or {}).get("channel") | |
| rows = [[("✅ Yes", f"{PREFIX}:y:{target}:{sent_id}:{channel or 0}"), ("❌ No", f"{PREFIX}:n")]] | |
| await self.client.send_message( | |
| job.chat_id, offer_text(job.new_name, episode_of(job.new_name), channel), buttons=panels.to_buttons(rows), | |
| parse_mode="html", link_preview=False, reply_to=sent_id if target == job.chat_id else job.msg_id, | |
| ) | |
| async def on_callback(self, event: Any, parts: list) -> None: | |
| action = parts[1] if len(parts) > 1 else "" | |
| if action == "n": | |
| await event.answer() | |
| await event.edit("👍 Okay, it is not uploaded to the Database Channel.", buttons=None, parse_mode="html") | |
| return | |
| try: | |
| if action != "y" or len(parts) != 5: | |
| raise ValueError(parts) | |
| source_chat, source_id, channel = int(parts[2]), int(parts[3]), int(parts[4]) | |
| except ValueError: | |
| await event.answer() | |
| return | |
| key = (event.chat_id, event.message_id) | |
| if key in self._busy: | |
| await event.answer("I'm already on it…") | |
| return | |
| self._busy.add(key) | |
| try: | |
| await event.answer() | |
| try: | |
| await event.edit(processing_text(""), buttons=None, parse_mode="html") | |
| except errors.MessageNotModifiedError: | |
| pass | |
| except BaseException: | |
| self._busy.discard(key) | |
| raise | |
| self._spawn(self._run(key, source_chat, source_id, channel or None)) | |
| # --------------------------------------------------------------- the upload | |
| async def _run(self, key: Tuple[int, int], source_chat: int, source_id: int, channel: Optional[int]) -> None: | |
| chat_id, msg_id = key | |
| flow: Optional[Flow] = None | |
| name = "" | |
| try: | |
| msg = await self.client.get_messages(source_chat, ids=source_id) | |
| if msg is None or getattr(msg, "file", None) is None: | |
| await self._show( | |
| chat_id, msg_id, "❌ <b>That episode is not there anymore</b>\nIt was deleted from the chat, so there is nothing to upload." | |
| ) | |
| return | |
| name = getattr(msg.file, "name", None) or "" | |
| episode = episode_of(name) | |
| # ---- 1. the Database Channel | |
| try: | |
| copy = await self._copy(msg, source_chat) | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as exc: # noqa: BLE001 | |
| log.warning("could not upload %s to the Database Channel: %s: %s", name, type(exc).__name__, exc) | |
| retry = [[("🔁 Try again", f"{PREFIX}:y:{source_chat}:{source_id}:{channel or 0}"), ("✖ Close", "x:close")]] | |
| await self._show( | |
| chat_id, msg_id, | |
| f"❌ <b>The upload to the Database Channel failed</b>\n{_name(name)}\n{_why(exc)}\n" | |
| "Check that the bot is an admin of the Database Channel and can post there.", | |
| retry, | |
| ) | |
| return | |
| if channel is None: | |
| await self._show( | |
| chat_id, msg_id, | |
| f"✅ <b>Uploaded to the Database Channel</b>\n{_name(name)}\nThis format has no Channel, so no post is made.", | |
| ) | |
| return | |
| # ---- 2. the communication chat | |
| flow = Flow(ref=secrets.token_hex(3)) | |
| self.flows[flow.ref] = flow # before the notice goes out, so that no answer can come back too early | |
| uploaded = "✅ Uploaded to the Database Channel" | |
| await self._show(chat_id, msg_id, processing_text(name, uploaded, "⏳ Waiting for the channel post bot…")) | |
| try: | |
| note = await self.client.send_message( | |
| self.cfg.comm_chat, | |
| notice_text(episode=episode, channel=channel, file=name, post_id=_id_of(copy), ref=flow.ref), | |
| parse_mode=None, link_preview=False, | |
| ) | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as exc: # noqa: BLE001 | |
| log.warning("could not write in the communication chat: %s: %s", type(exc).__name__, exc) | |
| await self._show( | |
| chat_id, msg_id, | |
| f"⚠️ <b>Uploaded, but I could not write in the communication chat</b>\n{_name(name)}\n{_why(exc)}\n" | |
| "Nothing was published. Is the bot a member of that chat?", | |
| ) | |
| return | |
| flow.notice_id = _id_of(note) or 0 | |
| # ---- 3. what the channel post bot says | |
| verdict, reply = await self._answers(flow, chat_id, msg_id, name, uploaded) | |
| if verdict == "ok": | |
| text = published_text(name, reply, channel) | |
| elif verdict == "fail": | |
| text = failed_text(name, classify(reply)[1]) | |
| elif verdict == "down": | |
| text = down_text(name, self.ack_wait) | |
| else: | |
| text = silent_text(name, self.result_wait) | |
| await self._show(chat_id, msg_id, text) | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as exc: # noqa: BLE001 | |
| log.exception("the upload to the Database Channel went wrong") | |
| await self._show(chat_id, msg_id, f"❌ <b>Something went wrong</b>\n{_name(name)}\n{_why(exc)}") | |
| finally: | |
| self._busy.discard(key) | |
| if flow is not None: | |
| self.flows.pop(flow.ref, None) | |
| async def _copy(self, msg: Any, source_chat: int) -> Any: | |
| """The episode in the Database Channel, copied on Telegram's side: forwarded without the author, so nothing is downloaded or | |
| uploaded. When Telegram refuses that (a chat that forbids forwarding), the same file is sent again by its id, which is as quick.""" | |
| try: | |
| return await self.client.forward_messages(self.cfg.db_channel, msg.id, from_peer=source_chat, drop_author=True) | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as exc: # noqa: BLE001 | |
| log.info("forwarding to the Database Channel was refused (%s: %s); sending the file by its id", type(exc).__name__, exc) | |
| return await self.client.send_file( | |
| self.cfg.db_channel, msg.media, caption=getattr(msg, "message", "") or "", | |
| formatting_entities=getattr(msg, "entities", None) or None, parse_mode=None, | |
| ) | |
| async def _answers(self, flow: Flow, chat_id: int, msg_id: int, name: str, uploaded: str) -> Tuple[str, str]: | |
| """Wait for the channel post bot: ``("ok" | "fail", its message)``, ``("down", "")`` when it never acknowledged the notice, | |
| ``("silent", "")`` when it acknowledged it but never said how it went.""" | |
| loop = asyncio.get_running_loop() | |
| acked = False | |
| deadline = loop.time() + self.ack_wait | |
| while True: | |
| try: | |
| kind, text = await asyncio.wait_for(flow.inbox.get(), timeout=max(0.01, deadline - loop.time())) | |
| except asyncio.TimeoutError: | |
| return ("silent" if acked else "down"), "" | |
| if kind != "ack": | |
| return kind, text | |
| if not acked: | |
| acked = True | |
| deadline = loop.time() + self.result_wait | |
| await self._show(chat_id, msg_id, processing_text(name, uploaded, "📥 The channel post bot got it and is publishing…")) | |
| async def _show(self, chat_id: int, msg_id: int, text: str, rows: Optional[list] = None) -> None: | |
| try: | |
| await self.client.edit_message( | |
| chat_id, msg_id, text, buttons=panels.to_buttons(rows) if rows else None, parse_mode="html", link_preview=False | |
| ) | |
| except errors.MessageNotModifiedError: | |
| pass | |
| except Exception as exc: # noqa: BLE001 | |
| log.warning("could not update the upload message: %s: %s", type(exc).__name__, exc) | |
| # ------------------------------------------------- the channel post bot's answers | |
| async def on_chat_message(self, event: Any) -> None: | |
| """A message of the channel post bot in the communication chat.""" | |
| try: | |
| if getattr(event, "chat_id", None) != self.cfg.comm_chat or getattr(event, "sender_id", None) != self.cfg.post_bot_id: | |
| return # (the event filter does this already) | |
| text = getattr(event, "raw_text", None) or "" | |
| kind = classify(text)[0] | |
| if kind == "other": | |
| return | |
| flow = self._flow_for(event, text) | |
| if flow is not None: | |
| flow.inbox.put_nowait((kind, text)) | |
| except Exception: # noqa: BLE001 | |
| log.exception("could not read a message in the communication chat") | |
| def _flow_for(self, event: Any, text: str) -> Optional[Flow]: | |
| """The upload an answer belongs to: the one whose Ref it repeats, or whose notice it replies to.""" | |
| ref = line_of("ref", text) | |
| if ref and ref in self.flows: | |
| return self.flows[ref] | |
| target = reply_target(event) | |
| if target: | |
| return next((f for f in self.flows.values() if f.notice_id == target), None) | |
| if len(self.flows) == 1: # a plain message, and only one upload waiting for an answer | |
| return next(iter(self.flows.values())) | |
| return None | |
| # ---------------------------------------------------------------- plumbing | |
| def _spawn(self, coro: Any) -> "asyncio.Future[Any]": | |
| task = asyncio.ensure_future(coro) | |
| self._tasks.add(task) | |
| task.add_done_callback(self._tasks.discard) | |
| return task | |
| async def idle(self) -> None: | |
| """Wait until no upload is under way (the tests use this).""" | |
| while self._tasks: | |
| await asyncio.gather(*list(self._tasks), return_exceptions=True) | |
| async def shutdown(self) -> None: | |
| for task in list(self._tasks): | |
| task.cancel() | |
| await asyncio.gather(*list(self._tasks), return_exceptions=True) | |
| self._tasks.clear() | |
| def _id_of(sent: Any) -> Optional[int]: | |
| if isinstance(sent, (list, tuple)): | |
| sent = sent[0] if sent else None | |
| value = getattr(sent, "id", None) | |
| return value if isinstance(value, int) else None | |