"""Filename and folder-name rules: sanitising, extension detection, collisions.""" from __future__ import annotations import os import re import unicodedata from pathlib import Path from errors import PermanentError VIDEO_EXTENSIONS = frozenset( {".mp4", ".m4v", ".mov", ".mkv", ".webm", ".avi", ".flv", ".wmv", ".ts", ".m2ts", ".mts", ".3gp", ".mpg", ".mpeg", ".ogv"} ) # Extensions that share a container format; a title saying ".mov" for an MP4 file is fine. _FAMILIES = ( frozenset({".mp4", ".m4v", ".mov", ".3gp"}), frozenset({".mkv", ".webm"}), frozenset({".ts", ".m2ts", ".mts"}), frozenset({".mpg", ".mpeg"}), ) MIME_TO_EXTENSION = { "video/mp4": ".mp4", "video/x-m4v": ".m4v", "video/quicktime": ".mov", "video/x-matroska": ".mkv", "video/webm": ".webm", "video/x-msvideo": ".avi", "video/x-flv": ".flv", "video/x-ms-wmv": ".wmv", "video/mp2t": ".ts", "video/3gpp": ".3gp", "video/mpeg": ".mpg", "video/ogg": ".ogv", } EXTENSION_TO_MIME = {ext: mime for mime, ext in MIME_TO_EXTENSION.items()} | { ".mpeg": "video/mpeg", ".m2ts": "video/mp2t", ".mts": "video/mp2t", } MAX_NAME_LENGTH = 200 _ALWAYS_REPLACE = re.compile(r"[/\\]") _STRICT_REPLACE = re.compile(r'[<>:"|?*]') def normalize_key(text: str) -> str: """Comparison key for folder names: NFC, collapsed whitespace, case-insensitive.""" return " ".join(unicodedata.normalize("NFC", text).split()).casefold() def clean_component(text: str | None, *, strict: bool = False) -> str: """Make a caption value safe as a Drive file/folder name while keeping Unicode. Control characters and path separators are replaced; letters, bullets, dots, parentheses and other punctuation are kept. ``strict`` additionally replaces characters Windows cannot store (useful with Google Drive for desktop). """ if not text: return "" text = unicodedata.normalize("NFC", text) text = "".join(" " if unicodedata.category(ch) == "Cc" else ch for ch in text) text = _ALWAYS_REPLACE.sub("_", text) if strict: text = _STRICT_REPLACE.sub("_", text).rstrip(". ") text = " ".join(text.split()) return truncate_name(text) def split_video_extension(name: str) -> tuple[str, str]: """Split off a *video* extension only, so 'Lecture 1.5' keeps its '.5'.""" stem, ext = os.path.splitext(name) if ext.lower() in VIDEO_EXTENSIONS and stem: return stem, ext return name, "" def truncate_name(name: str, limit: int = MAX_NAME_LENGTH) -> str: if len(name) <= limit: return name stem, ext = split_video_extension(name) return stem[: limit - len(ext)].rstrip() + ext def same_family(a: str, b: str) -> bool: a, b = a.lower(), b.lower() if a == b: return True return any(a in family and b in family for family in _FAMILIES) def extension_from_mime(mime_type: str | None) -> str | None: return MIME_TO_EXTENSION.get((mime_type or "").lower()) def mime_for_name(name: str, fallback: str | None = None) -> str: _, ext = split_video_extension(name) return EXTENSION_TO_MIME.get(ext.lower()) or fallback or "application/octet-stream" def detect_video_extension(path: Path) -> str | None: """Identify the container from magic bytes; ``None`` if unknown.""" with open(path, "rb") as fh: head = fh.read(4096) if len(head) >= 12 and head[4:8] == b"ftyp": brand = head[8:12] if brand == b"qt ": return ".mov" if brand.startswith(b"3g"): return ".3gp" if brand.startswith(b"M4V"): return ".m4v" return ".mp4" if head.startswith(b"\x1a\x45\xdf\xa3"): return ".webm" if b"webm" in head[:64] else ".mkv" if head[:4] == b"RIFF" and head[8:12] == b"AVI ": return ".avi" if head[:3] == b"FLV": return ".flv" if head[:8] == b"\x30\x26\xb2\x75\x8e\x66\xcf\x11": return ".wmv" if head[:4] in (b"\x00\x00\x01\xba", b"\x00\x00\x01\xb3"): return ".mpg" if head[:4] == b"OggS": return ".ogv" if len(head) > 376 and head[0] == 0x47 and head[188] == 0x47 and head[376] == 0x47: return ".ts" return None def build_filename( title: str | None, *, message_id: int, original_filename: str | None, mime_type: str | None, detected_ext: str | None, missing_title_policy: str = "fallback", mismatch_policy: str = "keep", strict: bool = False, ) -> tuple[str, list[str]]: """Return the Drive filename for a video plus any warnings worth recording. The Title field is used verbatim (sanitised). The extension is only added when the title has none; a title whose extension disagrees with the real container is kept as-is unless EXTENSION_MISMATCH_POLICY=append. """ warnings: list[str] = [] name = clean_component(title, strict=strict) if not name: if missing_title_policy == "fail": raise PermanentError( f"message {message_id} has no usable Title", public="Caption has no usable Title (MISSING_TITLE_POLICY=fail)" ) name = clean_component(original_filename, strict=strict) if name: warnings.append("Title missing; used the Telegram file name") else: name = f"telegram_{message_id}" warnings.append("Title missing; used a generated name") _, ext = split_video_extension(name) _, original_ext = split_video_extension(original_filename or "") actual = detected_ext or extension_from_mime(mime_type) or (original_ext.lower() or None) if not ext: if actual: name = truncate_name(name, MAX_NAME_LENGTH - len(actual)) + actual else: warnings.append("could not determine the file extension") elif detected_ext and not same_family(ext, detected_ext): warnings.append(f"title extension {ext} does not match detected container {detected_ext}") if mismatch_policy == "append": name = truncate_name(name, MAX_NAME_LENGTH - len(detected_ext)) + detected_ext return name, warnings def add_suffix(name: str, suffix: str) -> str: """'Class 01.mp4' + '(2)' -> 'Class 01 (2).mp4'.""" stem, ext = split_video_extension(name) if not ext: stem, ext = os.path.splitext(name) if len(ext) > 6 or not stem: stem, ext = name, "" return truncate_name(f"{stem} {suffix}{ext}")