Spaces:
Running
Running
| """The rename pipeline: download -> (ffmpeg) -> upload, or for a plain rename of a big file a copy straight from the download into the upload.""" | |
| from __future__ import annotations | |
| import logging | |
| import mimetypes | |
| import shutil | |
| import time | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple | |
| from telethon import errors, types | |
| from . import errorlog | |
| from . import fasttransfer as transfer | |
| from . import media | |
| from .config import Config | |
| from .db import Database | |
| from .errors import Cancelled, MediaError | |
| from .formats import apply_format | |
| from .parser import build_values, parse_info, render_caption, split_ext | |
| from .progress import ProgressReporter | |
| from .queue import Job | |
| from .utils import clip, esc, human_duration, human_size | |
| log = logging.getLogger(__name__) | |
| MIME_OVERRIDES = { | |
| ".mkv": "video/x-matroska", ".mka": "audio/x-matroska", ".webm": "video/webm", ".mp4": "video/mp4", | |
| ".m4v": "video/x-m4v", ".mov": "video/quicktime", ".avi": "video/x-msvideo", ".ts": "video/mp2t", | |
| ".mp3": "audio/mpeg", ".m4a": "audio/mp4", ".flac": "audio/flac", ".ogg": "audio/ogg", ".opus": "audio/ogg", | |
| ".wav": "audio/x-wav", ".srt": "application/x-subrip", ".ass": "text/x-ssa", ".zip": "application/zip", | |
| } | |
| # containers worth running ffprobe on | |
| PROBE_EXTS = media.VIDEO_CONTAINERS | media.AUDIO_TAG_EXTS | {".avi", ".ts", ".flv", ".wmv", ".mpg", ".mpeg", ".3gp", ".wav", ".aac", ".wma"} | |
| def guess_mime(name: str, fallback: Optional[str]) -> str: | |
| ext = split_ext(name)[1].lower() | |
| return MIME_OVERRIDES.get(ext) or mimetypes.guess_type(name)[0] or fallback or "application/octet-stream" | |
| def decide_send_type(setting: str, message: Any) -> str: | |
| """How the renamed file is sent: ``document`` | ``video`` | ``audio``.""" | |
| if setting in ("document", "video", "audio"): | |
| return setting | |
| if getattr(message, "video", None) is not None: | |
| return "video" | |
| if getattr(message, "audio", None) is not None: | |
| return "audio" | |
| return "document" | |
| def source_ext(message: Any, original_name: str) -> str: | |
| ext = split_ext(original_name)[1] | |
| if not ext: | |
| ext = getattr(message.file, "ext", None) or "" | |
| return ext | |
| def original_name_of(message: Any) -> str: | |
| f = message.file | |
| if f is not None and f.name: | |
| return f.name | |
| kind = "video" if getattr(message, "video", None) is not None else "audio" if getattr(message, "audio", None) is not None else "file" | |
| return f"{kind}_{message.id}{(f.ext if f is not None else '') or ''}" | |
| def make_caption(template: str, *, new_name: str, orig_name: str, ext: str, size: int, duration: float) -> Optional[str]: | |
| """The caption of a renamed file: the user's (or the format's) template with its variables filled in.""" | |
| parsed = parse_info(split_ext(orig_name)[0]) | |
| values = build_values(parsed, split_ext(orig_name)[0], ext) | |
| values.update( | |
| filename=split_ext(new_name)[0], original=orig_name, ext=ext.lstrip("."), | |
| size=human_size(size), duration=human_duration(duration) if duration else "", | |
| ) | |
| return render_caption(template, values)[:1000] or None | |
| def source_attributes(message: Any) -> Tuple[Optional[Any], Optional[Any]]: | |
| """The video and the audio details Telegram keeps with a file: ``(DocumentAttributeVideo, DocumentAttributeAudio)``, ``None`` for a missing one.""" | |
| attributes = getattr(getattr(message, "document", None), "attributes", None) or [] | |
| video = next((a for a in attributes if isinstance(a, types.DocumentAttributeVideo) and not a.round_message), None) | |
| audio = next((a for a in attributes if isinstance(a, types.DocumentAttributeAudio) and not a.voice), None) | |
| return video, audio | |
| async def embedded_thumb(client: Any, msg: Any, allow_stripped: bool = False) -> Optional[bytes]: | |
| """The preview Telegram keeps with a file. The tiny blurred one (a "stripped" size) only counts when ``allow_stripped``.""" | |
| thumbs = getattr(getattr(msg, "document", None), "thumbs", None) or [] | |
| if not any(not isinstance(t, types.PhotoStrippedSize) for t in thumbs) and not (allow_stripped and thumbs): | |
| return None | |
| try: | |
| return await client.download_media(msg, file=bytes, thumb=-1) or None | |
| except Exception as exc: # noqa: BLE001 | |
| log.debug("could not download the embedded thumbnail: %s", exc) | |
| return None | |
| class StreamPlan: | |
| """What a streamed rename sends along with the file, worked out without ever reading the file.""" | |
| attrs: List[Any] | |
| as_document: bool | |
| thumb: Optional[bytes] | |
| duration: float | |
| class Renamer: | |
| def __init__(self, client: Any, cfg: Config, db: Database): | |
| self.client = client | |
| self.cfg = cfg | |
| self.db = db | |
| # called as ``await on_delivered(job, chat, sent)`` once a file renamed with a manual format is in its chat (the offer | |
| # to upload it to the Database Channel, see dbflow); ``None``: nothing to do | |
| self.on_delivered: Optional[Callable[[Job, int, Any], Awaitable[None]]] = None | |
| # ------------------------------------------------------------------ | |
| async def _fresh_message(self, job: Job) -> Any: | |
| """Re-fetch the message so its file reference is valid even after a long queue wait.""" | |
| msg = await self.client.get_messages(job.chat_id, ids=job.msg_id) | |
| if msg is None or msg.file is None: | |
| raise RuntimeError("the original message was deleted") | |
| return msg | |
| def _check_space(self, work: Path, size: int, will_remux: bool) -> None: | |
| need = int(size * (2.2 if will_remux else 1.1)) + 64 * 1024 * 1024 | |
| free = shutil.disk_usage(work).free | |
| if free < need: | |
| raise RuntimeError(f"not enough free disk space ({human_size(free)} free, about {human_size(need)} needed)") | |
| # ------------------------------------------------------------------ | |
| async def process(self, job: Job) -> None: | |
| s = apply_format(await self.db.get_settings(job.user_id), job.fmt) # no format: the built-in template and caption | |
| work = Path(self.cfg.temp_dir) / f"job{job.id}" | |
| work.mkdir(parents=True, exist_ok=True) | |
| reporter = ProgressReporter(self.client, job.status_msg, job.id, job.new_name, self.cfg.progress_interval) | |
| started = time.monotonic() | |
| context = errorlog.error_context.set( # the log channel report says which file this was | |
| f"user {job.user_id} · job #{job.id}\n{job.file_name} → {job.new_name} ({human_size(job.size)})" | |
| ) | |
| try: | |
| await self._run(job, s, work, reporter, started) | |
| except Cancelled: | |
| job.state = "cancelled" | |
| await reporter.finish("❌ <b>Cancelled</b>\n<code>" + esc(clip(job.new_name, 120)) + "</code>") | |
| except Exception as exc: # noqa: BLE001 | |
| job.state = "failed" | |
| log.exception("Error processing file:") | |
| reason = clip(f"{type(exc).__name__}: {exc}", 300) | |
| await reporter.finish(f"❌ <b>Failed</b>\n<code>{esc(clip(job.new_name, 120))}</code>\n\n{esc(reason)}") | |
| finally: | |
| errorlog.error_context.reset(context) | |
| shutil.rmtree(work, ignore_errors=True) | |
| async def _custom_thumb(self, job: Job) -> Optional[bytes]: | |
| """The thumbnail the user chose: the one of the format being used (``None``: a frame of the video is used).""" | |
| return await self.db.get_format_thumb(job.user_id, job.fmt["id"]) if job.fmt else None | |
| async def _stream_plan( | |
| self, job: Job, s: Dict[str, Any], msg: Any, ext: str, send_type: str, size: int | |
| ) -> Optional[StreamPlan]: | |
| """What a streamed rename would send, or ``None`` when the file has to be looked at first (it is then saved and uploaded as usual). | |
| Streaming never saves the file, so nothing can read it: the video and audio details come from the ones Telegram keeps | |
| with the original, the thumbnail from the user's setting or the original's own preview. A rename that changes the | |
| file itself (metadata and tracks) cannot be streamed either. | |
| """ | |
| if not (self.cfg.fast_transfer and self.cfg.stream_transfer): | |
| return None | |
| video, audio = source_attributes(msg) | |
| source = video if send_type == "video" else audio if send_type == "audio" else (video or audio) | |
| duration = float(getattr(source, "duration", 0) or 0) | |
| reason = None | |
| if getattr(msg, "document", None) is None: | |
| reason = "it is not a document" | |
| elif size <= transfer.BIG_FILE_THRESHOLD: | |
| reason = "it is small, so there is little to gain" | |
| elif s["metadata_on"]: | |
| reason = "metadata and tracks change the file" | |
| elif send_type == "video" and video is None: | |
| reason = "Telegram has no video details to reuse" | |
| elif send_type == "audio" and audio is None: | |
| reason = "Telegram has no audio details to reuse" | |
| elif "{duration}" in (s["caption"] or "") and not duration: | |
| reason = "the caption wants the duration and only the file can tell" | |
| thumb = None | |
| if reason is None: | |
| thumb = await self._custom_thumb(job) | |
| if thumb is None and ext.lower() in media.VIDEO_CONTAINERS: # the usual way makes one from a frame of the file | |
| thumb = await embedded_thumb(self.client, msg) | |
| if thumb is None: | |
| reason = "there is no thumbnail to reuse" | |
| if reason is not None: | |
| log.info("%s is not streamed: %s", job.new_name, reason) | |
| return None | |
| attrs: List[Any] = [types.DocumentAttributeFilename(job.new_name)] | |
| if send_type == "video": | |
| attrs.append(types.DocumentAttributeVideo(duration=video.duration or 0, w=video.w or 1, h=video.h or 1, supports_streaming=True)) | |
| elif send_type == "audio": | |
| attrs.append( | |
| types.DocumentAttributeAudio( | |
| duration=int(audio.duration or 0), title=audio.title or None, performer=audio.performer or None | |
| ) | |
| ) | |
| return StreamPlan(attrs=attrs, as_document=(send_type == "document"), thumb=thumb, duration=duration) | |
| async def _stream( | |
| self, job: Job, s: Dict[str, Any], msg: Any, orig_name: str, ext: str, send_type: str, size: int, | |
| reporter: ProgressReporter, started: float, | |
| ) -> bool: | |
| """Rename by copying the file from the download straight into the upload: nothing is saved, so the whole takes about as long as | |
| the slower of the two. ``True`` when the file was sent; ``False`` when it cannot be streamed or streaming failed (the usual way follows).""" | |
| plan = await self._stream_plan(job, s, msg, ext, send_type, size) | |
| if plan is None: | |
| return False | |
| reporter.set_stage("⬇️⬆️ Streaming", size) | |
| began = time.monotonic() | |
| try: | |
| input_file = await transfer.relay( | |
| self.client, msg.document, job.new_name, connections=self.cfg.connections, progress=reporter, cancel=job.cancel | |
| ) | |
| except Cancelled: | |
| raise | |
| except Exception as exc: # noqa: BLE001 | |
| self._stream_failed(job, exc) | |
| return False | |
| seconds = time.monotonic() - began | |
| if job.cancel.is_set(): | |
| raise Cancelled() | |
| rates = f" · ⬇⬆ {human_size(size / seconds)}/s streamed" if seconds >= 1 else "" | |
| caption = make_caption(s["caption"], new_name=job.new_name, orig_name=orig_name, ext=ext, size=size, duration=plan.duration) | |
| try: | |
| await self._deliver( | |
| job, s, msg, input_file, caption=caption, attrs=plan.attrs, as_document=plan.as_document, thumb=plan.thumb, | |
| final_size=size, warnings=list(job.notes), rates=rates, started=started, reporter=reporter, | |
| ) | |
| except Exception as exc: # noqa: BLE001 | |
| if not transfer.is_file_part_error(exc): # only Telegram not accepting the parts is a streaming problem | |
| raise | |
| self._stream_failed(job, exc) | |
| return False | |
| transfer.stats.streamed += 1 | |
| return True | |
| def _stream_failed(self, job: Job, exc: BaseException) -> None: | |
| log.error("streaming %s failed (%s: %s); doing it the usual way", job.new_name, type(exc).__name__, exc, exc_info=exc) | |
| transfer.stats.stream_failed(exc) | |
| job.notes.append("Streaming didn't work for this file, so it was saved first and then uploaded.") | |
| async def _deliver( | |
| self, job: Job, s: Dict[str, Any], msg: Any, input_file: Any, *, caption: Optional[str], attrs: List[Any], | |
| as_document: bool, thumb: Optional[bytes], final_size: int, warnings: List[str], rates: str, started: float, | |
| reporter: ProgressReporter, | |
| ) -> None: | |
| """Send the uploaded file to its chat, then count it and say it is done.""" | |
| target = s.get("target_channel") or job.chat_id | |
| kwargs: Dict[str, Any] = dict( | |
| caption=caption, attributes=attrs, mime_type=guess_mime(job.new_name, getattr(msg.file, "mime_type", None)), | |
| force_document=as_document, thumb=thumb, parse_mode="html", | |
| ) | |
| if target == job.chat_id: | |
| kwargs["reply_to"] = job.msg_id | |
| try: | |
| sent = await self.client.send_file(target, input_file, **kwargs) | |
| except (errors.EntityBoundsInvalidError, errors.MediaCaptionTooLongError, errors.MessageTooLongError): | |
| # the caption template produced something Telegram rejects: send the file without a caption | |
| kwargs.update(caption=None) | |
| sent = await self.client.send_file(target, input_file, **kwargs) | |
| warnings.append("The caption was rejected by Telegram (check your caption template), so it was left out.") | |
| await self.db.record_job(job.user_id, final_size) | |
| job.state = "done" | |
| took = time.monotonic() - started | |
| note = "".join(f"\n⚠️ {esc(w)}" for w in warnings) | |
| where = "" if target == job.chat_id else "\n📣 Sent to your channel." | |
| await reporter.finish( | |
| f"✅ <b>Done</b> in {human_duration(took)}\n<code>{esc(clip(job.new_name, 120))}</code>\n{human_size(final_size)}{rates}{where}{note}" | |
| ) | |
| await self._after_delivery(job, target, sent) | |
| async def _after_delivery(self, job: Job, target: int, sent: Any) -> None: | |
| """A file renamed with a manual format may go on to the Database Channel: the offer comes after the file itself.""" | |
| if self.on_delivered is None or job.fmt is None or job.fmt.get("kind") != "manual": | |
| return | |
| try: | |
| await self.on_delivered(job, target, sent) | |
| except Exception: # noqa: BLE001 - the rename is done whatever happens here | |
| log.exception("could not offer the upload to the Database Channel") | |
| async def _run(self, job: Job, s: Dict[str, Any], work: Path, reporter: ProgressReporter, started: float) -> None: | |
| client, cfg = self.client, self.cfg | |
| msg = await self._fresh_message(job) | |
| orig_name = original_name_of(msg) | |
| ext = source_ext(msg, orig_name) | |
| send_type = decide_send_type(s["file_type"], msg) | |
| size = int(msg.file.size or job.size or 0) | |
| src = work / f"src{ext}" | |
| if await self._stream(job, s, msg, orig_name, ext, send_type, size, reporter, started): | |
| return | |
| self._check_space(work, size, will_remux=bool(s["metadata_on"])) | |
| # ---- download ------------------------------------------------------------ | |
| reporter.set_stage("⬇️ Downloading", size) | |
| began = time.monotonic() | |
| await transfer.download( | |
| client, msg, src, fast=cfg.fast_transfer, connections=cfg.connections, progress=reporter, cancel=job.cancel | |
| ) | |
| down_seconds = time.monotonic() - began | |
| if job.cancel.is_set(): | |
| raise Cancelled() | |
| # ---- probe / ffmpeg ----------------------------------------------------------- | |
| info: Optional[media.MediaInfo] = None | |
| current = src | |
| if ext.lower() in PROBE_EXTS or send_type in ("video", "audio"): | |
| try: | |
| info = await media.probe(src, job.cancel) | |
| except MediaError as exc: | |
| log.info("ffprobe could not read %s: %s", orig_name, exc) | |
| warnings = list(job.notes) | |
| plan = media.build_plan(info, s, ext) if (info is not None and s["metadata_on"]) else None | |
| if s["metadata_on"] and plan is None and info is not None and ext.lower() not in (media.VIDEO_CONTAINERS | media.AUDIO_TAG_EXTS): | |
| warnings.append(f"Metadata is not supported for {ext or 'this file type'}, so it was skipped.") | |
| if plan is not None: | |
| warnings.extend(plan.warnings) | |
| reporter.set_stage("⚙️ Applying metadata & tracks", 0) | |
| out = work / f"out{ext}" | |
| await media.apply_plan(src, out, plan, job.cancel) | |
| src.unlink(missing_ok=True) | |
| current = out | |
| info = await media.probe(current, job.cancel) | |
| final_size = current.stat().st_size | |
| # ---- how it will be sent ------------------------------------------------------- | |
| if send_type == "video" and (info is None or not info.of("video")): | |
| send_type = "document" | |
| if send_type == "audio" and (info is None or not info.of("audio")): | |
| send_type = "document" | |
| meta = s["meta"] if s["metadata_on"] else {} | |
| attrs = [types.DocumentAttributeFilename(job.new_name)] | |
| if send_type == "video": | |
| attrs.append( | |
| types.DocumentAttributeVideo( | |
| duration=info.duration or 0, w=info.width or 1, h=info.height or 1, supports_streaming=True | |
| ) | |
| ) | |
| elif send_type == "audio": | |
| attrs.append( | |
| types.DocumentAttributeAudio( | |
| duration=int(info.duration or 0), | |
| title=(meta.get("title") or info.tags.get("title") or None), | |
| performer=(meta.get("artist") or meta.get("author") or info.tags.get("artist") or None), | |
| ) | |
| ) | |
| thumb = await self._custom_thumb(job) | |
| if thumb is None and ext.lower() in media.VIDEO_CONTAINERS and info is not None and info.of("video"): | |
| thumb_path = work / "thumb.jpg" | |
| if await media.make_video_thumbnail(current, thumb_path, info.duration, job.cancel): | |
| thumb = thumb_path.read_bytes() | |
| # ---- caption --------------------------------------------------------------------- | |
| caption = make_caption( | |
| s["caption"], new_name=job.new_name, orig_name=orig_name, ext=ext, size=final_size, | |
| duration=(info.duration if info else 0) or 0, | |
| ) | |
| # ---- upload --------------------------------------------------------------------------- | |
| reporter.set_stage("⬆️ Uploading", final_size) | |
| began = time.monotonic() | |
| input_file = await transfer.upload( | |
| client, current, job.new_name, fast=cfg.fast_transfer, connections=cfg.connections, | |
| progress=reporter, cancel=job.cancel, | |
| ) | |
| up_seconds = time.monotonic() - began | |
| if job.cancel.is_set(): | |
| raise Cancelled() | |
| rates = "".join( # how fast each way went; a transfer under a second says nothing | |
| f" · {arrow} {human_size(amount / seconds)}/s" | |
| for arrow, amount, seconds in (("⬇", size, down_seconds), ("⬆", final_size, up_seconds)) | |
| if seconds >= 1 and amount > 0 | |
| ) | |
| await self._deliver( | |
| job, s, msg, input_file, caption=caption, attrs=attrs, as_document=(send_type == "document"), thumb=thumb, | |
| final_size=final_size, warnings=warnings, rates=rates, started=started, reporter=reporter, | |
| ) | |