"""I/O SRT: dataclass `Cue`, parse/write, ngân sách ký tự theo thời lượng. Port sạch từ bản cũ — không phụ thuộc gì ngoài stdlib để chạy được local ở Bước 1. """ from __future__ import annotations import re from dataclasses import dataclass, field from pathlib import Path from typing import Iterable, List, Optional # 00:01:02,500 --> 00:01:05,000 _TIME_RE = re.compile( r"(?P\d{1,2}):(?P\d{2}):(?P\d{2})[.,](?P\d{1,3})" ) _ARROW_RE = re.compile(r"-->") @dataclass class Cue: """Một dòng phụ đề.""" index: int start: float # giây end: float # giây text: str # văn bản gốc (ZH) speaker: Optional[str] = None # gán ở Bước 2 (diarization) translation: Optional[str] = None # bản dịch VN (Bước 4) note: Optional[str] = None # ghi chú xưng hô của model scene_id: Optional[int] = None # gán ở Bước 1 (gom cảnh) @property def duration(self) -> float: return max(0.0, self.end - self.start) @property def out_text(self) -> str: """Văn bản dùng khi ghi ra: ưu tiên bản dịch, fallback về gốc.""" return self.translation if self.translation else self.text def parse_time(token: str) -> float: """'00:01:02,500' -> giây (float).""" m = _TIME_RE.search(token) if not m: raise ValueError(f"Không đọc được mốc thời gian: {token!r}") h = int(m.group("h")) mi = int(m.group("m")) s = int(m.group("s")) ms = int(m.group("ms").ljust(3, "0")) # '5' -> 500 return h * 3600 + mi * 60 + s + ms / 1000.0 def seconds_to_srt_time(seconds: float) -> str: """giây -> '00:01:02,500'.""" if seconds < 0: seconds = 0.0 total_ms = int(round(seconds * 1000.0)) ms = total_ms % 1000 total_s = total_ms // 1000 s = total_s % 60 total_m = total_s // 60 m = total_m % 60 h = total_m // 60 return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}" def _normalize_text(raw_lines: List[str]) -> str: """Gộp các dòng văn bản của một cue, bỏ thẻ định dạng đơn giản.""" text = "\n".join(line.rstrip() for line in raw_lines).strip() return text def parse_srt(path: str | Path, encoding: str = "utf-8-sig") -> List[Cue]: """Đọc file SRT -> danh sách Cue. Bền với index lệch, dòng trống dư.""" p = Path(path) content = p.read_text(encoding=encoding, errors="replace") return parse_srt_text(content) def parse_srt_text(content: str) -> List[Cue]: content = content.replace("\r\n", "\n").replace("\r", "\n") blocks = re.split(r"\n\s*\n", content.strip()) cues: List[Cue] = [] auto_index = 0 for block in blocks: lines = [ln for ln in block.split("\n")] if not lines: continue # bỏ dòng trống đầu while lines and not lines[0].strip(): lines.pop(0) if not lines: continue idx_line = lines[0].strip() time_line_pos = 0 explicit_index: Optional[int] = None if idx_line.isdigit() and len(lines) >= 2 and _ARROW_RE.search(lines[1]): explicit_index = int(idx_line) time_line_pos = 1 elif _ARROW_RE.search(lines[0]): time_line_pos = 0 else: # block không hợp lệ (vd metadata) -> bỏ qua continue time_line = lines[time_line_pos] parts = _ARROW_RE.split(time_line) if len(parts) != 2: continue try: start = parse_time(parts[0]) end = parse_time(parts[1]) except ValueError: continue text = _normalize_text(lines[time_line_pos + 1 :]) auto_index += 1 cues.append( Cue( index=explicit_index if explicit_index is not None else auto_index, start=start, end=end, text=text, ) ) return cues def write_srt(cues: Iterable[Cue], path: str | Path, encoding: str = "utf-8") -> None: """Ghi danh sách Cue ra file SRT (dùng `out_text`).""" p = Path(path) p.parent.mkdir(parents=True, exist_ok=True) p.write_text(render_srt(cues), encoding=encoding) def render_srt(cues: Iterable[Cue]) -> str: chunks: List[str] = [] for i, cue in enumerate(cues, start=1): chunks.append( f"{i}\n" f"{seconds_to_srt_time(cue.start)} --> {seconds_to_srt_time(cue.end)}\n" f"{cue.out_text}".rstrip() ) return "\n\n".join(chunks) + "\n" def subtitle_char_budget( duration: float, chars_per_sec: float = 22.0, max_line_chars: int = 52, max_lines: int = 2, ) -> int: """Số ký tự tối đa nên hiển thị cho một cue, theo thời lượng và giới hạn dòng. Lấy min giữa (thời lượng × tốc độ đọc) và (số ký tự tối đa của các dòng). """ by_time = int(duration * chars_per_sec) by_lines = max_line_chars * max_lines # tối thiểu cho 1 dòng ngắn để không trả về 0 với cue siêu ngắn return max(by_lines if by_time <= 0 else min(by_time, by_lines), max_line_chars)