xiaoxia-saas/packages/domain/subtitle.py

"""字幕领域模型 — 带时间轴的字幕片段。"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List


@dataclass
class SubtitleWord:
    """单个词级别的字幕单元，带精确时间戳。"""

    text: str
    start: float  # 秒
    end: float  # 秒

    @property
    def duration(self) -> float:
        return max(0.0, self.end - self.start)


@dataclass
class SubtitleSegment:
    """一段字幕（一句话），带时间轴和词级信息。"""

    text: str
    start: float  # 秒
    end: float  # 秒
    words: List[SubtitleWord] = field(default_factory=list)

    @property
    def duration(self) -> float:
        return max(0.0, self.end - self.start)

    @property
    def char_count(self) -> int:
        return len(self.text)


@dataclass
class SubtitleTimeline:
    """完整的字幕时间轴，由多个片段组成。"""

    segments: List[SubtitleSegment] = field(default_factory=list)
    language: str = "zh"  # zh / en / ja 等
    total_duration: float = 0.0  # 音频总时长（秒）

    @property
    def segment_count(self) -> int:
        return len(self.segments)

    @property
    def total_chars(self) -> int:
        return sum(s.char_count for s in self.segments)

    def merge_short_segments(self, min_chars: int = 8) -> SubtitleTimeline:
        """合并过短的字幕片段，避免字幕跳动太快。"""
        if len(self.segments) <= 1:
            return self

        merged: List[SubtitleSegment] = []
        buffer: List[SubtitleSegment] = []

        for seg in self.segments:
            buffer.append(seg)
            total_chars = sum(s.char_count for s in buffer)
            if total_chars >= min_chars:
                merged.append(self._merge_segments(buffer))
                buffer = []

        # 剩余的合并到最后一个或单独成段
        if buffer:
            if merged and sum(s.char_count for s in buffer) < min_chars:
                # 太少了，合并到上一段
                last = merged.pop()
                merged.append(self._merge_segments([last] + buffer))
            else:
                merged.append(self._merge_segments(buffer))

        return SubtitleTimeline(
            segments=merged,
            language=self.language,
            total_duration=self.total_duration,
        )

    def split_long_segments(self, max_chars: int = 20) -> SubtitleTimeline:
        """拆分过长的字幕片段，按语义断句。"""
        new_segments: List[SubtitleSegment] = []

        for seg in self.segments:
            if seg.char_count <= max_chars:
                new_segments.append(seg)
                continue

            # 按标点符号拆分
            parts = self._split_text_by_punctuation(seg.text, max_chars)
            if len(parts) == 1:
                new_segments.append(seg)
                continue

            # 按字数比例分配时间
            total_chars = seg.char_count
            current_time = seg.start
            word_idx = 0
            all_words = seg.words.copy()

            for part in parts:
                part_chars = len(part)
                part_duration = seg.duration * (part_chars / total_chars)
                part_end = min(current_time + part_duration, seg.end)

                # 收集对应时间段的词
                part_words = []
                while word_idx < len(all_words) and all_words[word_idx].start < part_end:
                    part_words.append(all_words[word_idx])
                    word_idx += 1

                new_segments.append(
                    SubtitleSegment(
                        text=part,
                        start=current_time,
                        end=part_end,
                        words=part_words,
                    )
                )
                current_time = part_end

        return SubtitleTimeline(
            segments=new_segments,
            language=self.language,
            total_duration=self.total_duration,
        )

    @staticmethod
    def _merge_segments(segments: List[SubtitleSegment]) -> SubtitleSegment:
        if not segments:
            return SubtitleSegment(text="", start=0, end=0)
        return SubtitleSegment(
            text="".join(s.text for s in segments),
            start=segments[0].start,
            end=segments[-1].end,
            words=[w for s in segments for w in s.words],
        )

    @staticmethod
    def _split_text_by_punctuation(text: str, max_chars: int) -> List[str]:
        """按标点符号智能拆分长文本。"""
        # 中文常见句末标点
        sentence_end = "。！？!?"
        clause_pause = "，；：,;:"

        parts: List[str] = []
        current = ""

        for char in text:
            current += char

            if len(current) >= max_chars:
                # 超过长度，找最近的标点断开
                break_idx = -1
                for i in range(len(current) - 1, -1, -1):
                    if current[i] in sentence_end or current[i] in clause_pause:
                        break_idx = i + 1
                        break

                if break_idx > 0:
                    parts.append(current[:break_idx])
                    current = current[break_idx:]
                else:
                    # 没有标点，硬切
                    parts.append(current[:max_chars])
                    current = current[max_chars:]

            elif char in sentence_end:
                # 句末标点，如果长度够就断开
                if len(current) >= max_chars // 2:
                    parts.append(current)
                    current = ""

        if current:
            parts.append(current)

        return parts